@Test
  public void replicaThreadLeaderBasedReplicationTokenCatchUpForStandbyToLeaderTest() throws Exception {

    /*
      Setup:
      we have 3 nodes that have replicas belonging to same partitions:
      a) localNode (local node that hosts partitions)
      b) remoteNodeInLocalDC (remote node in local data center that shares the partitions)
      c) remoteNodeInRemoteDC (remote node in remote data center that shares the partitions)

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

    Map<DataNodeId, MockHost> hosts = new HashMap<>();
    hosts.put(remoteNodeInLocalDC, remoteHostInLocalDC);
    hosts.put(remoteNodeInRemoteDC, remoteHostInRemoteDC);
    int batchSize = 5;

    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 10 PUT messages in each of the partitions in remote nodes of local DC and remote DC that needs to be
    replicated at local node. They will be fetched in 2 cycles by having batch size of replication as 5.

    After first cycle, we will trigger partition state change for local standby to leader. We should see that new leader
    has updated token info, i.e. it should start fetching from message number 6.
     */

    // 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 + batchSize);
    }

    // Choose partitions that are leaders on both local and remote nodes
    Set<ReplicaId> leaderReplicasOnLocalAndRemoteNodes = new HashSet<>();

    // Track a standby replica which has leader partition on remote node. We will update the state of replica to leader after one cycle of replication
    // and verify that replication resumes from remote token.
    MockReplicaId localStandbyReplicaWithLeaderPartitionOnRemoteNode = null;

    List<? extends ReplicaId> localReplicas = clusterMap.getReplicaIds(replicationManager.dataNodeId);
    List<? extends ReplicaId> remoteReplicas = clusterMap.getReplicaIds(remoteNodeInRemoteDC);
    for (int i = 0; i < localReplicas.size(); i++) {
      MockReplicaId localReplica = (MockReplicaId) localReplicas.get(i);
      MockReplicaId remoteReplica = (MockReplicaId) remoteReplicas.get(i);
      if (localReplica.getReplicaState() == ReplicaState.LEADER
          && remoteReplica.getReplicaState() == ReplicaState.LEADER) {
        leaderReplicasOnLocalAndRemoteNodes.add(remoteReplicas.get(i));
      }
      if (localReplica.getReplicaState() == ReplicaState.STANDBY
          && remoteReplica.getReplicaState() == ReplicaState.LEADER
          && localStandbyReplicaWithLeaderPartitionOnRemoteNode == null) {
        localStandbyReplicaWithLeaderPartitionOnRemoteNode = localReplica;
      }
    }

    // replicate with remote node in remote DC
    crossColoReplicaThread.replicate();

    // 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 (RemoteReplicaInfo remoteReplicaInfo : remoteReplicaInfosForRemoteDC) {
      if (leaderReplicasOnLocalAndRemoteNodes.contains(remoteReplicaInfo.getReplicaId())) {
        assertEquals("remote token mismatch for leader replicas",
            ((MockFindToken) remoteReplicaInfo.getToken()).getIndex(), batchSize);
      } else {
        assertEquals("remote token should not move forward for standby replicas until missing keys are fetched",
            ((MockFindToken) remoteReplicaInfo.getToken()).getIndex(), 0);
      }
    }

    //Replicate with remote node in local dc
    intraColoReplicaThread.replicate();

    // verify that remote token will be moved for all intra-colo replicas with token index = batchSize-1
    for (RemoteReplicaInfo replicaInfo : remoteReplicaInfosForLocalDC) {
      assertEquals("mismatch in remote token set for intra colo replicas",
          ((MockFindToken) replicaInfo.getToken()).getIndex(), batchSize);
    }

    // process missing keys for cross colo replicas from previous metadata exchange
    for (RemoteReplicaInfo remoteReplicaInfo : remoteReplicaInfosForRemoteDC) {
      crossColoReplicaThread.processMissingKeysFromPreviousMetadataResponse(remoteReplicaInfo);
    }

    // verify that remote token will be moved for all (leader and standby) inter-colo replicas with token index = batchSize-1
    // as missing keys must now be obtained via intra-dc replication
    for (RemoteReplicaInfo replicaInfo : remoteReplicaInfosForRemoteDC) {
      assertEquals("mismatch in remote token set for inter colo replicas",
          ((MockFindToken) replicaInfo.getToken()).getIndex(), batchSize);
    }

    // If we have a local standby replica with leader partition on remote node, change its state to leader
    if (localStandbyReplicaWithLeaderPartitionOnRemoteNode != null) {
      MockPartitionId mockPartitionId =
          (MockPartitionId) localStandbyReplicaWithLeaderPartitionOnRemoteNode.getPartitionId();
      mockHelixParticipant.onPartitionBecomeLeaderFromStandby(mockPartitionId.toPathString());
    }

    // Trigger replication again with remote node in remote DC
    crossColoReplicaThread.replicate();

    // verify that remaining missing blobs of count batchSize have been fetched (token index = batchSize*2-2) for all leader replicas including new
    // leader localStandbyReplicaWithLeaderPartitionOnRemoteNode whose replication should be have resumed with remote token index = batchSize-1
    for (RemoteReplicaInfo remoteReplicaInfo : remoteReplicaInfosForRemoteDC) {
      if (leaderReplicasOnLocalAndRemoteNodes.contains(remoteReplicaInfo.getReplicaId())
          || (remoteReplicaInfo.getLocalReplicaId().equals(localStandbyReplicaWithLeaderPartitionOnRemoteNode))) {
        assertEquals("remote token mismatch for leader replicas",
            ((MockFindToken) remoteReplicaInfo.getToken()).getIndex(), batchSize * 2);
      } else {
        assertEquals("remote token should not move forward for standby replicas until missing keys are fetched",
            ((MockFindToken) remoteReplicaInfo.getToken()).getIndex(), batchSize);
      }
    }

    // Trigger replication again with remote node in local DC
    intraColoReplicaThread.replicate();

    // verify that remote token is moved forward for all intra-colo replicas.
    for (RemoteReplicaInfo replicaInfo : remoteReplicaInfosForLocalDC) {
      assertEquals("mismatch in remote token set for intra colo replicas",
          ((MockFindToken) replicaInfo.getToken()).getIndex(), batchSize * 2);
    }

    // process missing keys for cross colo replicas from previous metadata exchange
    for (RemoteReplicaInfo remoteReplicaInfo : remoteReplicaInfosForRemoteDC) {
      crossColoReplicaThread.processMissingKeysFromPreviousMetadataResponse(remoteReplicaInfo);
    }

    // verify that remote token is moved forward for all inter-colo replicas as missing keys must now be obtained
    // via intra-dc replication.
    for (RemoteReplicaInfo remoteReplicaInfo : remoteReplicaInfosForRemoteDC) {
      assertEquals("mismatch in remote token set for intra colo replicas",
          ((MockFindToken) remoteReplicaInfo.getToken()).getIndex(), batchSize * 2);
    }

    storageManager.shutdown();
  }
Test leader based replication to verify remote token is caught up for standby replicas and updated token is used
when their state transitions to leader.
@throws Exception
ambry-replication/src/test/java/com/github/ambry/replication/LeaderBasedReplicationTest.java