@Test
  public void partitionSelectionOnReplicaAddedOrRemovedTest() throws Exception {
    assumeTrue(!overrideEnabled);
    // create a HelixClusterManager with DynamicClusterChangeHandler
    Properties properties = new Properties();
    properties.putAll(props);
    ClusterMapConfig clusterMapConfig = new ClusterMapConfig(new VerifiableProperties(properties));
    HelixClusterManager helixClusterManager =
        new HelixClusterManager(clusterMapConfig, selfInstanceName, helixManagerFactory, new MetricRegistry());

    List<PartitionId> partitions = testPartitionLayout.getPartitionLayout().getPartitions(null);
    // partition1 has a replica to remove in local dc; partition2 has a replica to add in local dc
    // also add an extra replica from partition1 in remote dc
    PartitionId partition1 = partitions.get(0);
    PartitionId partition2 = partitions.get(partitions.size() - 1);
    int replicaCountForPart1 = partition1.getReplicaIds().size();
    int replicaCountForPart2 = partition2.getReplicaIds().size();
    List<DataNodeId> remoteNodesForPartition1 = getPartitionDataNodesFromDc(partition1, remoteDc);
    List<DataNodeId> localNodesForPartition2 = getPartitionDataNodesFromDc(partition2, localDc);
    // remove one replica from partition1 in local dc
    Replica replicaToRemove = (Replica) partition1.getReplicaIds()
        .stream()
        .filter(r -> r.getDataNodeId().getDatacenterName().equals(localDc))
        .findFirst()
        .get();
    testPartitionLayout.removeReplicaFromPartition(replicaToRemove);
    // add one replica of partition1 to remote dc
    Datacenter localDataCenter = null, remoteDataCenter = null;
    for (Datacenter datacenter : testHardwareLayout.getHardwareLayout().getDatacenters()) {
      if (datacenter.getName().equals(localDc)) {
        localDataCenter = datacenter;
      } else {
        remoteDataCenter = datacenter;
      }
    }
    assertNotNull("Remote data center is null", remoteDataCenter);
    assertNotNull("Local data center is null", localDataCenter);
    DataNode nodeToAddPartition1 =
        remoteDataCenter.getDataNodes().stream().filter(n -> !remoteNodesForPartition1.contains(n)).findFirst().get();
    testPartitionLayout.addReplicaToPartition(nodeToAddPartition1, (Partition) partition1);
    // add one replica of partition2 to node in local dc (the node should be different to currentNode as we didn't
    // populate bootstrap replica map for this test)
    DataNode nodeToAddPartition2 = localDataCenter.getDataNodes()
        .stream()
        .filter(n -> !localNodesForPartition2.contains(n) && n != currentNode)
        .findFirst()
        .get();
    testPartitionLayout.addReplicaToPartition(nodeToAddPartition2, (Partition) partition2);
    Utils.writeJsonObjectToFile(testPartitionLayout.getPartitionLayout().toJSONObject(), partitionLayoutPath);

    helixCluster.upgradeWithNewPartitionLayout(partitionLayoutPath,
        HelixBootstrapUpgradeUtil.HelixAdminOperation.BootstrapCluster);

    List<PartitionId> partitionsInManager = helixClusterManager.getAllPartitionIds(null);
    PartitionId ambryPartition1 =
        partitionsInManager.stream().filter(p -> p.toPathString().equals(partition1.toPathString())).findFirst().get();
    PartitionId ambryPartition2 =
        partitionsInManager.stream().filter(p -> p.toPathString().equals(partition2.toPathString())).findFirst().get();
    assertEquals("Replica count of partition1 is not expected", replicaCountForPart1,
        ambryPartition1.getReplicaIds().size());
    assertEquals("Replica count of partition2 is not expected", replicaCountForPart2 + 1,
        ambryPartition2.getReplicaIds().size());
    // Note that there are 3 partitions in total in cluster, partition[0] has 2 replicas, partition[1] has 3 replicas
    // and partition[2] has 4 replicas in local dc. Also the minimumLocalReplicaCount = 3, so this is to test if partition
    // selection helper is able to pick right partitions with local replica count >= 3.
    List<PartitionId> writablePartitions = helixClusterManager.getWritablePartitionIds(DEFAULT_PARTITION_CLASS);
    assertEquals("Number of writable partition is not expected", 2, writablePartitions.size());
    helixClusterManager.close();
  }
Test the case where replica is added or removed and {@link PartitionSelectionHelper} is able to incorporate cluster
map changes.
Test setup: (1) remove one replica of Partition1 from local dc;
            (2) add another replica of Partition2 to local dc;
            (3) add one replica of Partition1 to remote dc;
@throws Exception
ambry-clustermap/src/test/java/com/github/ambry/clustermap/ClusterChangeHandlerTest.java