mirror of
https://github.com/wassname/ray.git
synced 2026-07-19 11:27:32 +08:00
Fix ServiceBasedGcsGcsClientTest timing bug (#7365)
This commit is contained in:
@@ -272,6 +272,13 @@ class ServiceBasedGcsGcsClientTest : public RedisServiceManagerForTest {
|
||||
return status == std::future_status::ready;
|
||||
}
|
||||
|
||||
void WaitPendingDone(int ¤t_count, int expected_count) {
|
||||
auto condition = [¤t_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));
|
||||
|
||||
Reference in New Issue
Block a user