fix object manager ut

src/ray/object_manager/test/object_manager_stress_test.cc

@@ -33,10 +33,24 @@ namespace ray {
 
 using rpc::GcsNodeInfo;
 
-static inline void flushall_redis(void) {
+static inline bool flushall_redis(void) {
   redisContext *context = redisConnect("127.0.0.1", 6379);
+  if (context == nullptr || context->err) {
+    return false;
+  }
   freeReplyObject(redisCommand(context, "FLUSHALL"));
+  freeReplyObject(redisCommand(context, "SET NumRedisShards 1"));
+  freeReplyObject(redisCommand(context, "LPUSH RedisShards 127.0.0.1:6380"));
   redisFree(context);
+
+  redisContext *shard_context = redisConnect("127.0.0.1", 6380);
+  if (shard_context == nullptr || shard_context->err) {
+    return false;
+  }
+  freeReplyObject(redisCommand(shard_context, "FLUSHALL"));
+  redisFree(shard_context);
+
+  return true;
 }
 
 int64_t current_time_ms() {
@@ -72,6 +86,7 @@ class MockServer {
     node_info.set_object_manager_port(object_manager_port);
 
     ray::Status status = gcs_client_->Nodes().RegisterSelf(node_info, nullptr);
+    std::this_thread::sleep_for(std::chrono::milliseconds(5000));
     return status;
   }
 
@@ -86,7 +101,7 @@ class MockServer {
 class TestObjectManagerBase : public ::testing::Test {
  public:
   void SetUp() {
-    flushall_redis();
+    WaitForCondition(flushall_redis, 7000);
 
     // start store
     socket_name_1 = TestSetupUtil::StartObjectStore();
@@ -99,7 +114,7 @@ class TestObjectManagerBase : public ::testing::Test {
     // start first server
     gcs_server_socket_name_ = TestSetupUtil::StartGcsServer("127.0.0.1");
     gcs::GcsClientOptions client_options("127.0.0.1", 6379, /*password*/ "",
-                                         /*is_test_client=*/true);
+                                         /*is_test_client=*/false);
     gcs_client_1 = std::make_shared<gcs::ServiceBasedGcsClient>(client_options);
     RAY_CHECK_OK(gcs_client_1->Connect(main_service));
     ObjectManagerConfig om_config_1;
@@ -428,5 +443,6 @@ TEST_F(StressTestObjectManager, StartStressTestObjectManager) {
 int main(int argc, char **argv) {
   ::testing::InitGoogleTest(&argc, argv);
   ray::TEST_STORE_EXEC_PATH = std::string(argv[1]);
+  ray::TEST_GCS_SERVER_EXEC_PATH = std::string(argv[2]);
   return RUN_ALL_TESTS();
 }

src/ray/object_manager/test/object_manager_test.cc

@@ -39,6 +39,8 @@ using rpc::GcsNodeInfo;
 static inline void flushall_redis(void) {
   redisContext *context = redisConnect("127.0.0.1", 6379);
   freeReplyObject(redisCommand(context, "FLUSHALL"));
+  freeReplyObject(redisCommand(context, "SET NumRedisShards 1"));
+  freeReplyObject(redisCommand(context, "LPUSH RedisShards 127.0.0.1:6380"));
   redisFree(context);
 }
 
@@ -489,5 +491,6 @@ int main(int argc, char **argv) {
   ::testing::InitGoogleTest(&argc, argv);
   ray::TEST_STORE_EXEC_PATH = std::string(argv[1]);
   wait_timeout_ms = std::stoi(std::string(argv[2]));
+  ray::TEST_GCS_SERVER_EXEC_PATH = std::string(argv[3]);
   return RUN_ALL_TESTS();
 }

src/ray/test/run_object_manager_tests.sh

@@ -25,18 +25,22 @@ fi
 REDIS_MODULE="./bazel-bin/libray_redis_module.so"
 LOAD_MODULE_ARGS=(--loadmodule "${REDIS_MODULE}")
 STORE_EXEC="./bazel-bin/plasma_store_server"
+GCS_SERVER_EXEC="./bazel-bin/gcs_server"
 
 # Allow cleanup commands to fail.
 bazel run //:redis-cli -- -p 6379 shutdown || true
+bazel run //:redis-cli -- -p 6380 shutdown || true
 sleep 1s
 bazel run //:redis-server -- --loglevel warning "${LOAD_MODULE_ARGS[@]}" --port 6379 &
+bazel run //:redis-server -- --loglevel warning "${LOAD_MODULE_ARGS[@]}" --port 6380 &
 sleep 1s
 # Run tests.
-./bazel-bin/object_manager_stress_test $STORE_EXEC
+./bazel-bin/object_manager_stress_test $STORE_EXEC $GCS_SERVER_EXEC
 sleep 1s
 # Use timeout=1000ms for the Wait tests.
-./bazel-bin/object_manager_test $STORE_EXEC 1000
+./bazel-bin/object_manager_test $STORE_EXEC 1000 $GCS_SERVER_EXEC
 bazel run //:redis-cli -- -p 6379 shutdown
+bazel run //:redis-cli -- -p 6380 shutdown
 sleep 1s
 
 # Include raylet integration test once it's ready.