consensus_manager_test.cpp

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 */
 
#include "platform/networkstrate/consensus_manager.h"
 
#include <gtest/gtest.h>
 
#include <future>
 
#include "common/test/test_macros.h"
#include "platform/networkstrate/mock_replica_communicator.h"
 
namespace resdb {
namespace {
 
using ::resdb::testing::EqualsProto;
using ::testing::_;
using ::testing::AtLeast;
using ::testing::Invoke;
using ::testing::Pointee;
using ::testing::Return;
using ::testing::Test;
 
ReplicaInfo GenerateReplicaInfo(const std::string& ip, int port) {
  ReplicaInfo info;
  info.set_ip(ip);
  info.set_port(port);
  return info;
}
 
class MockConsensusManager : public ConsensusManager {
 public:
  MockConsensusManager(const ResDBConfig& config) : ConsensusManager(config) {}
 
  MOCK_METHOD(std::unique_ptr<ReplicaCommunicator>, GetReplicaClient,
              (const std::vector<ReplicaInfo>&, bool), (override));
  MOCK_METHOD(int, ConsensusCommit,
              (std::unique_ptr<Context>, std::unique_ptr<Request>), (override));
  MOCK_METHOD(std::vector<ReplicaInfo>, GetReplicas, (), (override));
  MOCK_METHOD(std::vector<ReplicaInfo>, GetClientReplicas, (), (override));
 
 public:
  void BroadCast(const Request& request) {
    return ConsensusManager::BroadCast(request);
  }
  void UpdateBroadCastClient() {
    return ConsensusManager::UpdateBroadCastClient();
  }
 
  ReplicaCommunicator* GetBroadCastClient() {
    return ConsensusManager::GetBroadCastClient();
  }
 
 public:
  int Dispatch(std::unique_ptr<Context> context,
               std::unique_ptr<Request> request) {
    return ConsensusManager::Dispatch(std::move(context), std::move(request));
  }
};
 
class ConsensusManagerTest : public Test {
 public:
  ConsensusManagerTest()
      : replicas_({GenerateReplicaInfo("127.0.0.1", 1234),
                   GenerateReplicaInfo("127.0.0.1", 1235)}),
        self_info_(GenerateReplicaInfo("127.0.0.1", 10000)),
        config_(replicas_, self_info_, KeyInfo(), CertificateInfo()) {
    ResConfigData data;
    auto region = data.add_region();
    region->set_region_id(1);
    for (ReplicaInfo info : config_.GetReplicaInfos()) {
      *region->add_replica_info() = info;
    }
    config_.SetConfigData(data);
 
    Stats::GetGlobalStats(1);
 
    impl_ = std::make_unique<MockConsensusManager>(config_);
 
    ON_CALL(*impl_, GetReplicas).WillByDefault(Return(replicas_));
  }
 
  ~ConsensusManagerTest() { impl_->Stop(); }
 
 protected:
  std::vector<ReplicaInfo> replicas_;
  ReplicaInfo self_info_;
  ResDBConfig config_;
  std::unique_ptr<MockConsensusManager> impl_;
};
 
TEST_F(ConsensusManagerTest, SendHB) {
  std::promise<bool> hb;
  std::future<bool> hb_done = hb.get_future();
  EXPECT_CALL(*impl_, GetReplicaClient)
      .Times(AtLeast(1))
      .WillRepeatedly(
          Invoke([&](const std::vector<ReplicaInfo>& replicas, bool) {
            auto client = std::make_unique<MockReplicaCommunicator>(replicas);
            EXPECT_CALL(*client, SendHeartBeat)
                .WillRepeatedly(Invoke([&](const Request& hb_info) {
                  hb.set_value(true);
                  return 0;
                }));
            return client;
          }));
  impl_->Start();
  hb_done.get();
}
 
TEST_F(ConsensusManagerTest, SendHBToClient) {
  std::promise<bool> hb;
  std::future<bool> hb_done = hb.get_future();
 
  EXPECT_CALL(*impl_, GetClientReplicas)
      .Times(AtLeast(1))
      .WillRepeatedly(Invoke([&]() {
        ReplicaInfo client_info;
        client_info.set_ip("127.0.0.1");
        client_info.set_port(54321);
        return std::vector<ReplicaInfo>({client_info});
      }));
 
  EXPECT_CALL(*impl_, GetReplicaClient)
      .Times(AtLeast(1))
      .WillRepeatedly(
          Invoke([&](const std::vector<ReplicaInfo>& replicas, bool) {
            EXPECT_EQ(replicas.size(), 3);
            bool find = false;
            for (auto& info : replicas) {
              if (info.ip() == "127.0.0.1" && info.port() == 54321) {
                find = true;
              }
            }
            EXPECT_TRUE(find);
            auto client = std::make_unique<MockReplicaCommunicator>(replicas);
            EXPECT_CALL(*client, SendHeartBeat)
                .WillRepeatedly(Invoke([&](const Request& hb_info) {
                  hb.set_value(true);
                  return 0;
                }));
            return client;
          }));
  impl_->Start();
  hb_done.get();
}
 
TEST_F(ConsensusManagerTest, SendHBToClientWithTworegion) {
  ResConfigData data;
  auto region = data.add_region();
  region->set_region_id(1);
  *region->add_replica_info() = GenerateReplicaInfo("127.0.0.1", 1234);
  *region->add_replica_info() = GenerateReplicaInfo("127.0.0.1", 1235);
 
  region = data.add_region();
  region->set_region_id(2);
  *region->add_replica_info() = GenerateReplicaInfo("127.0.0.1", 1234);
  *region->add_replica_info() = GenerateReplicaInfo("127.0.0.1", 1235);
  config_.SetConfigData(data);
  impl_ = std::make_unique<MockConsensusManager>(config_);
 
  std::promise<bool> hb;
  std::future<bool> hb_done = hb.get_future();
 
  EXPECT_CALL(*impl_, GetClientReplicas)
      .Times(AtLeast(1))
      .WillRepeatedly(Invoke([&]() {
        ReplicaInfo client_info;
        client_info.set_ip("127.0.0.1");
        client_info.set_port(54321);
        return std::vector<ReplicaInfo>({client_info});
      }));
 
  EXPECT_CALL(*impl_, GetReplicaClient)
      .Times(AtLeast(1))
      .WillRepeatedly(
          Invoke([&](const std::vector<ReplicaInfo>& replicas, bool) {
            EXPECT_EQ(replicas.size(), 5);
            bool find = false;
            for (auto& info : replicas) {
              if (info.ip() == "127.0.0.1" && info.port() == 54321) {
                find = true;
              }
            }
            EXPECT_TRUE(find);
            auto client = std::make_unique<MockReplicaCommunicator>(replicas);
            EXPECT_CALL(*client, SendHeartBeat)
                .WillRepeatedly(Invoke([&](const Request& hb_info) {
                  hb.set_value(true);
                  return 0;
                }));
            return client;
          }));
  impl_->Start();
  hb_done.get();
}
 
TEST_F(ConsensusManagerTest, BroadCast) {
  std::promise<bool> bc;
  std::future<bool> bc_done = bc.get_future();
  EXPECT_CALL(*impl_, GetReplicaClient)
      .WillRepeatedly(
          Invoke([&](const std::vector<ReplicaInfo>& replicas, bool long_conn) {
            EXPECT_EQ(replicas.size(), replicas_.size());
            auto client = std::make_unique<MockReplicaCommunicator>(replicas);
            EXPECT_CALL(*client, SendMessage(_))
                .WillRepeatedly(
                    Invoke([&](const google::protobuf::Message& request) {
                      try {
                        bc.set_value(true);
                      } catch (...) {
                      }
                      return 0;
                    }));
            return client;
          }));
 
  impl_->Start();
  impl_->UpdateBroadCastClient();
  impl_->BroadCast(Request());
  bc_done.get();
}
 
TEST_F(ConsensusManagerTest, DiscoverNewClient) {
  config_.SetSignatureVerifierEnabled(false);
  impl_ = std::make_unique<MockConsensusManager>(config_);
  impl_->Start();
 
  EXPECT_CALL(*impl_, GetReplicaClient)
      .WillRepeatedly(
          Invoke([&](const std::vector<ReplicaInfo>& replicas, bool long_conn) {
            EXPECT_EQ(replicas.size(), replicas_.size());
            auto client = std::make_unique<MockReplicaCommunicator>(replicas);
            return client;
          }));
 
  HeartBeatInfo hb;
  auto* public_key = hb.add_public_keys();
 
  public_key->mutable_public_key_info()->set_ip("127.0.0.1");
  public_key->mutable_public_key_info()->set_port(12345);
  public_key->mutable_public_key_info()->set_node_id(5);
  public_key->mutable_public_key_info()->set_type(CertificateKeyInfo::CLIENT);
 
  auto request = std::make_unique<Request>();
  request->set_type(Request::TYPE_HEART_BEAT);
  hb.SerializeToString(request->mutable_data());
 
  EXPECT_EQ(impl_->Dispatch(std::make_unique<Context>(), std::move(request)),
            0);
 
  ReplicaCommunicator* bc_client = impl_->GetBroadCastClient();
  EXPECT_NE(bc_client, nullptr);
 
  std::vector<ReplicaInfo> client_infos = bc_client->GetClientReplicas();
  EXPECT_EQ(client_infos.size(), 1);
}
 
TEST_F(ConsensusManagerTest, DispatchOK) {
  Request expected_request;
  expected_request.set_type(Request::TYPE_CLIENT_REQUEST);
  EXPECT_CALL(*impl_,
              ConsensusCommit(_, Pointee(EqualsProto(expected_request))))
      .Times(1);
 
  impl_->Start();
  {
    auto request = std::make_unique<Request>();
    request->set_type(Request::TYPE_HEART_BEAT);
    auto replica_info = GenerateReplicaInfo("127.0.0.1", 1234);
    replica_info.SerializeToString(request->mutable_data());
 
    EXPECT_EQ(impl_->Dispatch(std::make_unique<Context>(), std::move(request)),
              0);
  }
 
  {
    auto request = std::make_unique<Request>();
    request->set_type(Request::TYPE_CLIENT_REQUEST);
    EXPECT_EQ(impl_->Dispatch(std::make_unique<Context>(), std::move(request)),
              0);
  }
}
 
}  // namespace
 
}  // namespace resdb