@Test
  public void replicaThreadLeaderBasedReplicationForPUTMessagesTest() throws Exception {
    int batchSize = 4;

    Map<DataNodeId, MockHost> hosts = new HashMap<>();
    hosts.put(remoteNodeInLocalDC, remoteHostInLocalDC);
    hosts.put(remoteNodeInRemoteDC, remoteHostInRemoteDC);
    MockNetworkClientFactory mockNetworkClientFactory = new MockNetworkClientFactory(hosts, clusterMap, batchSize, new MockFindTokenHelper(new BlobIdFactory(clusterMap), replicationConfig));
    Pair<StorageManager, ReplicationManager> managers =
        createStorageManagerAndReplicationManager(clusterMap, clusterMapConfig, mockHelixParticipant,
            mockNetworkClientFactory);
    StorageManager storageManager = managers.getFirst();
    MockReplicationManager replicationManager = (MockReplicationManager) managers.getSecond();

    /*
    Scenario:

    we have 3 nodes that have replicas belonging to same partitions:
       a) local node
       b) remote node in local DC
       c) remote node in remote DC

    Each node have few of its partitions as leaders and others are standby. They are randomly assigned during creation
    of replicas for mock partitions.

    We have 4 PUT messages in each of the partitions in remote nodes of local DC and remote DC that needs to be
    replicated at local node.

    Steps:

    1. Replicate (send metadata exchange and get messages) with remote node in remote DC (cross-colo replication).

       Expectations:
        a) We should see that metadata exchange is sent for all replicas while GET messages are only sent for leader replicas.
        b) All the PUT messages should be replicated locally and remote token should be moved forward for leader partitions.
        c) For non-leader replicas, metadata response should be stored locally.

    2. Replicate (send metadata exchange and get messages) with remote node in local DC (intra-colo replication).

       Expectations:
        a) Metadata exchange and GET messages are sent for all replicas.
        b) PUT messages should be replicated locally for all replicas.
        c) Missing messages in stored metadata response of non-leader replicas for remoteNodeInRemoteDC should become empty
           and remote token should be advanced.

     */

    // set mock local stores on all remoteReplicaInfos which will used during replication.
    for (PartitionId partitionId : replicationManager.partitionToPartitionInfo.keySet()) {
      localHost.addStore(partitionId, null);
      Store localStore = localHost.getStore(partitionId);
      localStore.start();
      List<RemoteReplicaInfo> remoteReplicaInfos =
          replicationManager.partitionToPartitionInfo.get(partitionId).getRemoteReplicaInfos();
      remoteReplicaInfos.forEach(remoteReplicaInfo -> remoteReplicaInfo.setLocalStore(localStore));
    }

    // get remote replicas and replica thread for remote host on local datacenter
    ReplicaThread intraColoReplicaThread = replicationManager.dataNodeIdToReplicaThread.get(remoteNodeInLocalDC);
    List<RemoteReplicaInfo> remoteReplicaInfosForLocalDC =
        intraColoReplicaThread.getRemoteReplicaInfos().get(remoteNodeInLocalDC);

    // get remote replicas and replica thread for remote host on remote datacenter
    ReplicaThread crossColoReplicaThread = replicationManager.dataNodeIdToReplicaThread.get(remoteNodeInRemoteDC);
    List<RemoteReplicaInfo> remoteReplicaInfosForRemoteDC =
        crossColoReplicaThread.getRemoteReplicaInfos().get(remoteNodeInRemoteDC);

    // mock helix transition state from standby to leader for local leader partitions
    List<? extends ReplicaId> replicaIds = clusterMap.getReplicaIds(replicationManager.dataNodeId);
    for (ReplicaId replicaId : replicaIds) {
      MockReplicaId mockReplicaId = (MockReplicaId) replicaId;
      if (mockReplicaId.getReplicaState() == ReplicaState.LEADER) {
        MockPartitionId mockPartitionId = (MockPartitionId) replicaId.getPartitionId();
        mockHelixParticipant.onPartitionBecomeLeaderFromStandby(mockPartitionId.toPathString());
      }
    }

    //Add put messages to all partitions on remoteHost1 and remoteHost2
    List<PartitionId> partitionIds = clusterMap.getWritablePartitionIds(null);
    for (PartitionId partitionId : partitionIds) {
      // add batchSize messages to the remoteHost1 and remote host 2 from which local host will replicate.
      addPutMessagesToReplicasOfPartition(partitionId, Arrays.asList(remoteHostInLocalDC, remoteHostInRemoteDC),
          batchSize);
    }

    // Choose partitions that are leaders on both local and remote nodes
    Set<ReplicaId> leaderReplicasOnLocalAndRemoteNodes =
        getRemoteLeaderReplicasWithLeaderPartitionsOnLocalNode(clusterMap, replicationManager.dataNodeId,
            remoteNodeInRemoteDC);

    // Replicate with remoteHost2 in remote data center.
    crossColoReplicaThread.replicate();
    List<ReplicaThread.ExchangeMetadataResponse> responseListForRemoteNodeInRemoteDC =
        crossColoReplicaThread.getExchangeMetadataResponsesInEachCycle().get(remoteNodeInRemoteDC);

    // Metadata requests should be sent to both leader and standby replicas.
    assertEquals("Response should contain a response for each replica", remoteReplicaInfosForRemoteDC.size(),
        responseListForRemoteNodeInRemoteDC.size());

    // verify that missing messages size equals to the min{batch size, number of PUT messages} placed on remote hosts
    int expectedIndex = batchSize - 1;
    for (ReplicaThread.ExchangeMetadataResponse exchangeMetadataResponse : responseListForRemoteNodeInRemoteDC) {
      assertEquals("mismatch in number of missing messages", batchSize, exchangeMetadataResponse.missingStoreMessages.size());
    }

    // Filter leader replicas to fetch missing keys
    List<RemoteReplicaInfo> leaderReplicas = new ArrayList<>();
    List<ReplicaThread.ExchangeMetadataResponse> exchangeMetadataResponseListForLeaderReplicas = new ArrayList<>();
    crossColoReplicaThread.getLeaderReplicaList(remoteReplicaInfosForRemoteDC, responseListForRemoteNodeInRemoteDC,
        leaderReplicas, exchangeMetadataResponseListForLeaderReplicas);

    // verify that only leader replicas in remoteHost2 are chosen for fetching missing messages.
    Set<ReplicaId> remoteReplicasToFetchInReplicaThread =
        leaderReplicas.stream().map(RemoteReplicaInfo::getReplicaId).collect(Collectors.toSet());
    assertThat("mismatch in leader remote replicas to fetch missing keys", leaderReplicasOnLocalAndRemoteNodes, is(remoteReplicasToFetchInReplicaThread));

    // verify that the remote token will be moved for leader replicas and will remain 0 for standby replicas as
    // missing messages are not fetched yet.
    for (int i = 0; i < remoteReplicaInfosForRemoteDC.size(); i++) {
      if (leaderReplicasOnLocalAndRemoteNodes.contains(remoteReplicaInfosForRemoteDC.get(i).getReplicaId())) {
        assertEquals("remote token mismatch for leader replicas", remoteReplicaInfosForRemoteDC.get(i).getToken(),
            responseListForRemoteNodeInRemoteDC.get(i).remoteToken);
      } else {
        assertEquals("missing keys in metadata response should be stored for standby replicas",
            remoteReplicaInfosForRemoteDC.get(i).getExchangeMetadataResponse().missingStoreMessages.size(),
            responseListForRemoteNodeInRemoteDC.get(i).missingStoreMessages.size());
        assertThat("remote token should not move forward for standby replicas until missing keys are fetched",
            remoteReplicaInfosForRemoteDC.get(i).getToken(), not(responseListForRemoteNodeInRemoteDC.get(i).remoteToken));
      }
    }

    if (!leaderReplicas.isEmpty()) {
      // replicate again, nothing going to change
      crossColoReplicaThread.replicate();
      // Only the leader replicas should be able to replicate
      List<ReplicaThread.ExchangeMetadataResponse> responseListForRemoteNodeInRemoteDCNew =
          crossColoReplicaThread.getExchangeMetadataResponsesInEachCycle().get(remoteNodeInRemoteDC);
      assertEquals("Only Leader replicas would be able to replicate again", leaderReplicas.size(),
          responseListForRemoteNodeInRemoteDCNew.size());
      for (ReplicaThread.ExchangeMetadataResponse exchangeMetadataResponse : responseListForRemoteNodeInRemoteDCNew) {
        assertEquals("Should already caught up", 0, exchangeMetadataResponse.missingStoreMessages.size());
      }
    }

    // Replication with remoteHost1 in local data center
    intraColoReplicaThread.replicate();
    List<ReplicaThread.ExchangeMetadataResponse> responseForRemoteNodeInLocalDC =
        intraColoReplicaThread.getExchangeMetadataResponsesInEachCycle().get(remoteNodeInLocalDC);

    for (int i = 0; i < responseForRemoteNodeInLocalDC.size(); i++) {
      assertEquals("mismatch in remote token set for intra colo replicas",
          remoteReplicaInfosForLocalDC.get(i).getToken(), (responseForRemoteNodeInLocalDC.get(i).remoteToken));
    }

    // replicate again, so we can process missing keys for cross colo standby replicas from previous metadata exchange
    crossColoReplicaThread.replicate();

    // Remote token for all cross-colo replicas (leader and standby) should have moved forward now as the missing keys
    // for standbys are received via intra-dc replication.
    for (int i = 0; i < responseListForRemoteNodeInRemoteDC.size(); i++) {
      assertEquals("mismatch in remote token set for cross colo replicas", remoteReplicaInfosForRemoteDC.get(i).getToken(), (responseListForRemoteNodeInRemoteDC.get(i).remoteToken));
    }

    // verify replication metrics to track number of cross colo get requests and cross colo bytes fetch rate for standby
    // replicas should be 0 since all missing blobs are obtained via local leader in intra-dc replication.
    String remoteDataCenter = remoteReplicaInfosForRemoteDC.get(0).getReplicaId().getDataNodeId().getDatacenterName();
    assertEquals("mismatch in number of cross colo get requests tracked for standby replicas",
        crossColoReplicaThread.getReplicationMetrics().interColoReplicationGetRequestCountForStandbyReplicas.get(
            remoteDataCenter).getCount(), 0);
    assertEquals("mismatch in bytes fetch rate for cross colo get requests tracked for standby replicas",
        crossColoReplicaThread.getReplicationMetrics().interColoReplicationFetchBytesRateForStandbyReplicas.get(
            remoteDataCenter).getCount(), 0);

    storageManager.shutdown();
  }
Test leader based replication to ensure token is advanced correctly for standby replicas when missing PUT messages
are fetched via intra-dc replication.
@throws Exception
ambry-replication/src/test/java/com/github/ambry/replication/LeaderBasedReplicationTest.java