Fix ServiceBasedGcsGcsClientTest timing bug (#7365)

This commit is contained in:
micafan
2020-02-28 12:01:02 -06:00
committed by GitHub
parent 93fe4b0b58
commit 3f8b1d2756
@@ -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));