From bb95f655c4e6028d86090b9db083f499466dd345 Mon Sep 17 00:00:00 2001
From: Carlos Segarra <carlos@carlossegarra.com>
Date: Fri, 15 Oct 2021 15:10:05 +0000
Subject: [PATCH] pr comments

---
 tests/test/scheduler/test_exec_graph.cpp | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/tests/test/scheduler/test_exec_graph.cpp b/tests/test/scheduler/test_exec_graph.cpp
index db3b99cab..d40357063 100644
--- a/tests/test/scheduler/test_exec_graph.cpp
+++ b/tests/test/scheduler/test_exec_graph.cpp
@@ -77,6 +77,9 @@ TEST_CASE_METHOD(MpiBaseTestFixture,
     faabric::scheduler::MpiWorld world;
     msg.set_ismpi(true);
 
+    // Update the result for the master message
+    sch.setFunctionResult(msg);
+
     // Build the message vector to reconstruct the graph
     std::vector<faabric::Message> messages(worldSize);
     for (int rank = 0; rank < worldSize; rank++) {
@@ -91,11 +94,6 @@ TEST_CASE_METHOD(MpiBaseTestFixture,
 
     world.destroy();
 
-    // Update the result for the master message
-    sch.setFunctionResult(msg);
-    // Wait for the scheduler to set the result on the MPI non-master messages
-    SLEEP_MS(500);
-
     // Build expected graph
     ExecGraphNode nodeB1 = { .msg = messages.at(1) };
     ExecGraphNode nodeB2 = { .msg = messages.at(2) };
@@ -108,6 +106,7 @@ TEST_CASE_METHOD(MpiBaseTestFixture,
     ExecGraph expected{ .rootNode = nodeA };
 
     // Check the execution graph
+    sch.getFunctionResult(msg.id(), 500);
     ExecGraph actual = sch.getFunctionExecGraph(msg.id());
     REQUIRE(countExecGraphNodes(actual) == worldSize);
     REQUIRE(countExecGraphNodes(expected) == worldSize);