commit 098751016e93bf10c3dfcb8dadfdb15c7d15ad44
author Shen Li <cs.shenli@gmail.com> Thu Jan 07 19:43:44 2021 -0800
committer Facebook GitHub Bot <facebook-github-bot@users.noreply.github.com> Thu Jan 07 19:50:50 2021 -0800
tree 4b665ce1f7733d9ecfb25a6d460a44747c965b47
parent 1f795e1a9b937e5ca83e589030f69092d33ce4dc

Completely Remove FutureMessage from RPC cpp tests (#50027)

Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/50027

Test Plan: Imported from OSS

Reviewed By: lw

Differential Revision: D25753815

Pulled By: mrshenli

fbshipit-source-id: 85b9b03fec52b4175288ac3a401285607744b451

test/cpp/rpc/e2e_test_base.h

diff --git a/test/cpp/rpc/e2e_test_base.h b/test/cpp/rpc/e2e_test_base.h
index 61d9a73..cea5079 100644
--- a/test/cpp/rpc/e2e_test_base.h
+++ b/test/cpp/rpc/e2e_test_base.h
@@ -90,14 +90,15 @@
     ScriptCall scriptCall(op, {t1, t2, /* alpha */ 1});
 
     // Send the RPC and return result.
-    auto response = RpcAgent::toFutureMessage(
-        autograd::sendMessageWithAutograd(
-            *rpcAgent,
-            rpcAgent->getWorkerInfo("worker"),
-            std::move(scriptCall).toMessage()))
-        ->wait();
+    auto response = autograd::sendMessageWithAutograd(
+        *rpcAgent,
+        rpcAgent->getWorkerInfo("worker"),
+        std::move(scriptCall).toMessage());
+    response->waitAndThrow();
+
     MessageType messageType = MessageType::FORWARD_AUTOGRAD_RESP;
-    auto wrappedResponse = deserializeResponse(response, messageType);
+    auto wrappedResponse = deserializeResponse(
+        std::move(*response->value().toCustomClass<Message>()), messageType);
     return static_cast<ScriptResp&>(*wrappedResponse).value().toTensor();
   }