From 3f8b1d275611a7c21945e455bb7faf3bd33d7277 Mon Sep 17 00:00:00 2001
From: micafan <550435771@qq.com>
Date: Sat, 29 Feb 2020 02:01:02 +0800
Subject: [PATCH] Fix ServiceBasedGcsGcsClientTest timing bug (#7365)

---
 .../test/service_based_gcs_client_test.cc     | 21 +++++++++++++------
 1 file changed, 15 insertions(+), 6 deletions(-)

diff --git a/src/ray/gcs/gcs_client/test/service_based_gcs_client_test.cc b/src/ray/gcs/gcs_client/test/service_based_gcs_client_test.cc
index 168257610..e5f9166c3 100644
--- a/src/ray/gcs/gcs_client/test/service_based_gcs_client_test.cc
+++ b/src/ray/gcs/gcs_client/test/service_based_gcs_client_test.cc
@@ -272,6 +272,13 @@ class ServiceBasedGcsGcsClientTest : public RedisServiceManagerForTest {
     return status == std::future_status::ready;
   }
 
+  void WaitPendingDone(int &current_count, int expected_count) {
+    auto condition = [&current_count, expected_count]() {
+      return current_count == expected_count;
+    };
+    EXPECT_TRUE(WaitForCondition(condition, timeout_ms_.count()));
+  }
+
   std::shared_ptr<rpc::JobTableData> GenJobTableData(JobID job_id) {
     auto job_table_data = std::make_shared<rpc::JobTableData>();
     job_table_data->set_job_id(job_id.Binary());
@@ -482,6 +489,7 @@ TEST_F(ServiceBasedGcsGcsClientTest, TestNodeInfo) {
   ClientID node2_id = ClientID::FromRandom();
   auto gcs_node2_info = GenGcsNodeInfo(node2_id.Binary());
   ASSERT_TRUE(RegisterNode(gcs_node2_info));
+  WaitPendingDone(register_count, 2);
 
   // Get node list
   std::vector<rpc::GcsNodeInfo> node_list = GetNodeInfoList();
@@ -495,6 +503,8 @@ TEST_F(ServiceBasedGcsGcsClientTest, TestNodeInfo) {
 
   // Unregister node
   ASSERT_TRUE(UnregisterNode(node2_id));
+  WaitPendingDone(unregister_count, 2);
+
   node_list = GetNodeInfoList();
   EXPECT_EQ(node_list.size(), 2);
   EXPECT_EQ(node_list[0].state(),
@@ -527,15 +537,15 @@ TEST_F(ServiceBasedGcsGcsClientTest, TestNodeResources) {
   resource->set_resource_capacity(1.0);
   resource_map[key] = resource;
   ASSERT_TRUE(UpdateResources(node_id, resource_map));
+  WaitPendingDone(add_count, 1);
   auto get_resources_result = GetResources(node_id);
   ASSERT_TRUE(get_resources_result.count(key));
 
   // Delete resources
   ASSERT_TRUE(DeleteResources(node_id, {key}));
+  WaitPendingDone(remove_count, 1);
   get_resources_result = GetResources(node_id);
   ASSERT_TRUE(get_resources_result.empty());
-  EXPECT_EQ(add_count, 1);
-  EXPECT_EQ(remove_count, 1);
 }
 
 TEST_F(ServiceBasedGcsGcsClientTest, TestNodeHeartbeat) {
@@ -560,14 +570,13 @@ TEST_F(ServiceBasedGcsGcsClientTest, TestNodeHeartbeat) {
   auto heartbeat = std::make_shared<rpc::HeartbeatTableData>();
   heartbeat->set_client_id(node_id.Binary());
   ASSERT_TRUE(ReportHeartbeat(heartbeat));
+  WaitPendingDone(heartbeat_count, 1);
 
   // Report batch heartbeat
   auto batch_heartbeat = std::make_shared<rpc::HeartbeatBatchTableData>();
   batch_heartbeat->add_batch()->set_client_id(node_id.Binary());
   ASSERT_TRUE(ReportBatchHeartbeat(batch_heartbeat));
-
-  EXPECT_EQ(heartbeat_count, 1);
-  EXPECT_EQ(heartbeat_batch_count, 1);
+  WaitPendingDone(heartbeat_batch_count, 1);
 }
 
 TEST_F(ServiceBasedGcsGcsClientTest, TestTaskInfo) {
@@ -606,7 +615,7 @@ TEST_F(ServiceBasedGcsGcsClientTest, TestTaskInfo) {
   ClientID node_id = ClientID::FromRandom();
   auto task_lease = GenTaskLeaseData(task_id.Binary(), node_id.Binary());
   ASSERT_TRUE(AddTaskLease(task_lease));
-  EXPECT_EQ(task_lease_count, 2);
+  WaitPendingDone(task_lease_count, 2);
 
   RAY_CHECK_OK(gcs_client_->Tasks().AsyncUnsubscribeTaskLease(task_id, nullptr));
   ASSERT_TRUE(AddTaskLease(task_lease));