diff --git a/src/ray/common/test_util.cc b/src/ray/common/test_util.cc index a16434cba..bac0183d9 100644 --- a/src/ray/common/test_util.cc +++ b/src/ray/common/test_util.cc @@ -164,6 +164,11 @@ void TestSetupUtil::StopRaylet(const std::string &raylet_socket_name) { KillProcessBySocketName(raylet_socket_name); } +bool WaitReady(std::future future, const std::chrono::milliseconds &timeout_ms) { + auto status = future.wait_for(timeout_ms); + return status == std::future_status::ready && future.get(); +} + bool WaitForCondition(std::function condition, int timeout_ms) { int wait_time = 0; while (true) { diff --git a/src/ray/common/test_util.h b/src/ray/common/test_util.h index a35afb9a5..0077945f9 100644 --- a/src/ray/common/test_util.h +++ b/src/ray/common/test_util.h @@ -16,6 +16,7 @@ #include #include +#include #include #include "gtest/gtest.h" @@ -42,6 +43,13 @@ class RayObject; // Magic argument to signal to mock_worker we should check message order. static const int64_t SHOULD_CHECK_MESSAGE_ORDER = 123450000; +/// Wait until the future is ready, or timeout is reached. +/// +/// \param[in] future The future to wait for. +/// \param[in] timeout_ms Timeout in milliseconds to wait for for. +/// \return Whether the future is ready. +bool WaitReady(std::future future, const std::chrono::milliseconds &timeout_ms); + /// Wait until the condition is met, or timeout is reached. /// /// \param[in] condition The condition to wait for. diff --git a/src/ray/gcs/gcs_client/test/global_state_accessor_test.cc b/src/ray/gcs/gcs_client/test/global_state_accessor_test.cc index 1198c0908..bc5988fca 100644 --- a/src/ray/gcs/gcs_client/test/global_state_accessor_test.cc +++ b/src/ray/gcs/gcs_client/test/global_state_accessor_test.cc @@ -77,11 +77,6 @@ class GlobalStateAccessorTest : public ::testing::Test { TestSetupUtil::FlushAllRedisServers(); } - bool WaitReady(std::future future, const std::chrono::milliseconds &timeout_ms) { - auto status = future.wait_for(timeout_ms); - return status == std::future_status::ready && future.get(); - } - // GCS server. gcs::GcsServerConfig config; std::unique_ptr gcs_server_; 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 4fbeb1cf5..ab8c561ec 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 @@ -578,11 +578,6 @@ class ServiceBasedGcsClientTest : public ::testing::Test { return WaitReady(promise.get_future(), timeout_ms_); } - bool WaitReady(std::future future, const std::chrono::milliseconds &timeout_ms) { - auto status = future.wait_for(timeout_ms); - return status == std::future_status::ready && future.get(); - } - void CheckActorData(const gcs::ActorTableData &actor, rpc::ActorTableData_ActorState expected_state) { ASSERT_TRUE(actor.state() == expected_state); diff --git a/src/ray/gcs/gcs_server/test/gcs_server_rpc_test.cc b/src/ray/gcs/gcs_server/test/gcs_server_rpc_test.cc index 814f429b8..fd9e991b0 100644 --- a/src/ray/gcs/gcs_server/test/gcs_server_rpc_test.cc +++ b/src/ray/gcs/gcs_server/test/gcs_server_rpc_test.cc @@ -433,11 +433,6 @@ class GcsServerTest : public ::testing::Test { return WaitReady(promise.get_future(), timeout_ms_); } - bool WaitReady(const std::future &future, uint64_t timeout_ms) { - auto status = future.wait_for(std::chrono::milliseconds(timeout_ms)); - return status == std::future_status::ready; - } - protected: // Gcs server std::unique_ptr gcs_server_; @@ -449,7 +444,7 @@ class GcsServerTest : public ::testing::Test { std::unique_ptr client_call_manager_; // Timeout waiting for gcs server reply, default is 5s - const uint64_t timeout_ms_ = 5000; + const std::chrono::milliseconds timeout_ms_{5000}; }; TEST_F(GcsServerTest, TestActorInfo) { diff --git a/src/ray/gcs/pubsub/test/gcs_pub_sub_test.cc b/src/ray/gcs/pubsub/test/gcs_pub_sub_test.cc index 1f96daffb..55e7d0dd0 100644 --- a/src/ray/gcs/pubsub/test/gcs_pub_sub_test.cc +++ b/src/ray/gcs/pubsub/test/gcs_pub_sub_test.cc @@ -91,11 +91,6 @@ class GcsPubSubTest : public ::testing::Test { return WaitReady(promise.get_future(), timeout_ms_); } - bool WaitReady(std::future future, const std::chrono::milliseconds &timeout_ms) { - auto status = future.wait_for(timeout_ms); - return status == std::future_status::ready && future.get(); - } - template void WaitPendingDone(const std::vector &data, int expected_count) { auto condition = [this, &data, expected_count]() {