Revert "Revert "[test] add unit test for PR #17634 (#18585)" (#18830)"

BUILD.bazel

@@ -1513,6 +1513,21 @@ cc_test(
     ],
 )
 
+cc_test(
+    name = "gcs_actor_scheduler_mock_test",
+    size = "small",
+    srcs = [
+        "src/ray/gcs/gcs_server/test/gcs_actor_scheduler_mock_test.cc",
+    ],
+    copts = COPTS,
+    tags = ["team:core"],
+    deps = [
+        ":gcs_server_lib",
+        ":ray_mock",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 cc_test(
     name = "gcs_based_actor_scheduler_test",
     size = "small",

src/mock/ray/gcs/gcs_server/gcs_node_manager.h

@@ -17,6 +17,7 @@ namespace gcs {
 
 class MockGcsNodeManager : public GcsNodeManager {
  public:
+  MockGcsNodeManager() : GcsNodeManager(nullptr, nullptr) {}
   MOCK_METHOD(void, HandleRegisterNode,
               (const rpc::RegisterNodeRequest &request, rpc::RegisterNodeReply *reply,
                rpc::SendReplyCallback send_reply_callback),

src/mock/ray/gcs/pubsub/gcs_pub_sub.h

@@ -0,0 +1,27 @@
+// Copyright 2021 The Ray Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+namespace ray {
+namespace gcs {
+
+class MockGcsPubSub : public GcsPubSub {
+ public:
+  MOCK_METHOD(Status, Publish,
+              (const std::string &channel, const std::string &id, const std::string &data,
+               const StatusCallback &done),
+              (override));
+};
+
+}  // namespace gcs
+}  // namespace ray

src/mock/ray/gcs/store_client/in_memory_store_client.h

@@ -0,0 +1,66 @@
+// Copyright 2021 The Ray Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+namespace ray {
+namespace gcs {
+
+class MockInMemoryStoreClient : public InMemoryStoreClient {
+ public:
+  MOCK_METHOD(Status, AsyncPut,
+              (const std::string &table_name, const std::string &key,
+               const std::string &data, const StatusCallback &callback),
+              (override));
+  MOCK_METHOD(Status, AsyncPutWithIndex,
+              (const std::string &table_name, const std::string &key,
+               const std::string &index_key, const std::string &data,
+               const StatusCallback &callback),
+              (override));
+  MOCK_METHOD(Status, AsyncGet,
+              (const std::string &table_name, const std::string &key,
+               const OptionalItemCallback<std::string> &callback),
+              (override));
+  MOCK_METHOD(Status, AsyncGetByIndex,
+              (const std::string &table_name, const std::string &index_key,
+               (const MapCallback<std::string, std::string> &callback)),
+              (override));
+  MOCK_METHOD(Status, AsyncGetAll,
+              (const std::string &table_name,
+               (const MapCallback<std::string, std::string> &callback)),
+              (override));
+  MOCK_METHOD(Status, AsyncDelete,
+              (const std::string &table_name, const std::string &key,
+               const StatusCallback &callback),
+              (override));
+  MOCK_METHOD(Status, AsyncDeleteWithIndex,
+              (const std::string &table_name, const std::string &key,
+               const std::string &index_key, const StatusCallback &callback),
+              (override));
+  MOCK_METHOD(Status, AsyncBatchDelete,
+              (const std::string &table_name, const std::vector<std::string> &keys,
+               const StatusCallback &callback),
+              (override));
+  MOCK_METHOD(Status, AsyncBatchDeleteWithIndex,
+              (const std::string &table_name, const std::vector<std::string> &keys,
+               const std::vector<std::string> &index_keys,
+               const StatusCallback &callback),
+              (override));
+  MOCK_METHOD(Status, AsyncDeleteByIndex,
+              (const std::string &table_name, const std::string &index_key,
+               const StatusCallback &callback),
+              (override));
+  MOCK_METHOD(int, GetNextJobID, (), (override));
+};
+
+}  // namespace gcs
+}  // namespace ray

src/mock/ray/gcs/store_client/redis_store_client.h

@@ -0,0 +1,66 @@
+// Copyright 2021 The Ray Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+namespace ray {
+namespace gcs {
+
+class MockRedisStoreClient : public RedisStoreClient {
+ public:
+  MOCK_METHOD(Status, AsyncPut,
+              (const std::string &table_name, const std::string &key,
+               const std::string &data, const StatusCallback &callback),
+              (override));
+  MOCK_METHOD(Status, AsyncPutWithIndex,
+              (const std::string &table_name, const std::string &key,
+               const std::string &index_key, const std::string &data,
+               const StatusCallback &callback),
+              (override));
+  MOCK_METHOD(Status, AsyncGet,
+              (const std::string &table_name, const std::string &key,
+               const OptionalItemCallback<std::string> &callback),
+              (override));
+  MOCK_METHOD(Status, AsyncGetByIndex,
+              (const std::string &table_name, const std::string &index_key,
+               (const MapCallback<std::string, std::string> &callback)),
+              (override));
+  MOCK_METHOD(Status, AsyncGetAll,
+              (const std::string &table_name,
+               (const MapCallback<std::string, std::string> &callback)),
+              (override));
+  MOCK_METHOD(Status, AsyncDelete,
+              (const std::string &table_name, const std::string &key,
+               const StatusCallback &callback),
+              (override));
+  MOCK_METHOD(Status, AsyncDeleteWithIndex,
+              (const std::string &table_name, const std::string &key,
+               const std::string &index_key, const StatusCallback &callback),
+              (override));
+  MOCK_METHOD(Status, AsyncBatchDelete,
+              (const std::string &table_name, const std::vector<std::string> &keys,
+               const StatusCallback &callback),
+              (override));
+  MOCK_METHOD(Status, AsyncBatchDeleteWithIndex,
+              (const std::string &table_name, const std::vector<std::string> &keys,
+               const std::vector<std::string> &index_keys,
+               const StatusCallback &callback),
+              (override));
+  MOCK_METHOD(Status, AsyncDeleteByIndex,
+              (const std::string &table_name, const std::string &index_key,
+               const StatusCallback &callback),
+              (override));
+  MOCK_METHOD(int, GetNextJobID, (), (override));
+};
+
+}  // namespace gcs
+}  // namespace ray

src/mock/ray/gcs/store_client/store_client.h

@@ -0,0 +1,66 @@
+// Copyright 2021 The Ray Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+namespace ray {
+namespace gcs {
+
+class MockStoreClient : public StoreClient {
+ public:
+  MOCK_METHOD(Status, AsyncPut,
+              (const std::string &table_name, const std::string &key,
+               const std::string &data, const StatusCallback &callback),
+              (override));
+  MOCK_METHOD(Status, AsyncPutWithIndex,
+              (const std::string &table_name, const std::string &key,
+               const std::string &index_key, const std::string &data,
+               const StatusCallback &callback),
+              (override));
+  MOCK_METHOD(Status, AsyncGet,
+              (const std::string &table_name, const std::string &key,
+               const OptionalItemCallback<std::string> &callback),
+              (override));
+  MOCK_METHOD(Status, AsyncGetByIndex,
+              (const std::string &table_name, const std::string &index_key,
+               (const MapCallback<std::string, std::string> &callback)),
+              (override));
+  MOCK_METHOD(Status, AsyncGetAll,
+              (const std::string &table_name,
+               (const MapCallback<std::string, std::string> &callback)),
+              (override));
+  MOCK_METHOD(Status, AsyncDelete,
+              (const std::string &table_name, const std::string &key,
+               const StatusCallback &callback),
+              (override));
+  MOCK_METHOD(Status, AsyncDeleteWithIndex,
+              (const std::string &table_name, const std::string &key,
+               const std::string &index_key, const StatusCallback &callback),
+              (override));
+  MOCK_METHOD(Status, AsyncBatchDelete,
+              (const std::string &table_name, const std::vector<std::string> &keys,
+               const StatusCallback &callback),
+              (override));
+  MOCK_METHOD(Status, AsyncBatchDeleteWithIndex,
+              (const std::string &table_name, const std::vector<std::string> &keys,
+               const std::vector<std::string> &index_keys,
+               const StatusCallback &callback),
+              (override));
+  MOCK_METHOD(Status, AsyncDeleteByIndex,
+              (const std::string &table_name, const std::string &index_key,
+               const StatusCallback &callback),
+              (override));
+  MOCK_METHOD(int, GetNextJobID, (), (override));
+};
+
+}  // namespace gcs
+}  // namespace ray

src/mock/ray/pubsub/publisher.h

@@ -0,0 +1,100 @@
+// Copyright 2021 The Ray Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+namespace ray {
+namespace pubsub {
+namespace pub_internal {
+
+template <typename KeyIdType>
+class MockSubscriptionIndex : public SubscriptionIndex<KeyIdType> {
+ public:
+};
+
+}  // namespace pub_internal
+}  // namespace pubsub
+}  // namespace ray
+
+namespace ray {
+namespace pubsub {
+namespace pub_internal {
+
+class MockLongPollConnection : public LongPollConnection {
+ public:
+};
+
+}  // namespace pub_internal
+}  // namespace pubsub
+}  // namespace ray
+
+namespace ray {
+namespace pubsub {
+namespace pub_internal {
+
+class MockSubscriber : public Subscriber {
+ public:
+};
+
+}  // namespace pub_internal
+}  // namespace pubsub
+}  // namespace ray
+
+namespace ray {
+namespace pubsub {
+
+class MockPublisherInterface : public PublisherInterface {
+ public:
+  MOCK_METHOD(bool, RegisterSubscription,
+              (const rpc::ChannelType channel_type, const SubscriberID &subscriber_id,
+               const std::string &key_id_binary),
+              (override));
+  MOCK_METHOD(void, Publish,
+              (const rpc::ChannelType channel_type, const rpc::PubMessage &pub_message,
+               const std::string &key_id_binary),
+              (override));
+  MOCK_METHOD(void, PublishFailure,
+              (const rpc::ChannelType channel_type, const std::string &key_id_binary),
+              (override));
+  MOCK_METHOD(bool, UnregisterSubscription,
+              (const rpc::ChannelType channel_type, const SubscriberID &subscriber_id,
+               const std::string &key_id_binary),
+              (override));
+};
+
+}  // namespace pubsub
+}  // namespace ray
+
+namespace ray {
+namespace pubsub {
+
+class MockPublisher : public Publisher {
+ public:
+  MOCK_METHOD(bool, RegisterSubscription,
+              (const rpc::ChannelType channel_type, const SubscriberID &subscriber_id,
+               const std::string &key_id_binary),
+              (override));
+  MOCK_METHOD(void, Publish,
+              (const rpc::ChannelType channel_type, const rpc::PubMessage &pub_message,
+               const std::string &key_id_binary),
+              (override));
+  MOCK_METHOD(void, PublishFailure,
+              (const rpc::ChannelType channel_type, const std::string &key_id_binary),
+              (override));
+  MOCK_METHOD(bool, UnregisterSubscription,
+              (const rpc::ChannelType channel_type, const SubscriberID &subscriber_id,
+               const std::string &key_id_binary),
+              (override));
+};
+
+}  // namespace pubsub
+}  // namespace ray