@Test
  public void partitionSelectionHelperTest() {
    // set up partitions for tests
    // 2 partitions with 3 replicas in two datacenters "DC1" and "DC2" (class "max-replicas-all-sites")
    // 2 partitions with 3 replicas in "DC1" and 1 replica in "DC2" (class "max-local-one-remote")
    // 2 partitions with 3 replicas in "DC2" and 1 replica in "DC1" (class "max-local-one-remote")
    // minimum number of replicas required for choosing writable partition is 3.
    final String dc1 = "DC1";
    final String dc2 = "DC2";
    final String maxReplicasAllSites = "max-replicas-all-sites";
    final String maxLocalOneRemote = "max-local-one-remote";
    final int minimumLocalReplicaCount = 3;

    MockDataNodeId dc1Dn1 = getDataNodeId("dc1dn1", dc1);
    MockDataNodeId dc1Dn2 = getDataNodeId("dc1dn2", dc1);
    MockDataNodeId dc1Dn3 = getDataNodeId("dc1dn3", dc1);
    MockDataNodeId dc2Dn1 = getDataNodeId("dc2dn1", dc2);
    MockDataNodeId dc2Dn2 = getDataNodeId("dc2dn2", dc2);
    MockDataNodeId dc2Dn3 = getDataNodeId("dc2dn3", dc2);
    List<MockDataNodeId> allDataNodes = Arrays.asList(dc1Dn1, dc1Dn2, dc1Dn3, dc2Dn1, dc2Dn2, dc2Dn3);
    MockPartitionId everywhere1 = new MockPartitionId(1, maxReplicasAllSites, allDataNodes, 0);
    MockPartitionId everywhere2 = new MockPartitionId(2, maxReplicasAllSites, allDataNodes, 0);
    MockPartitionId majorDc11 =
        new MockPartitionId(3, maxLocalOneRemote, Arrays.asList(dc1Dn1, dc1Dn2, dc1Dn3, dc2Dn1), 0);
    MockPartitionId majorDc12 =
        new MockPartitionId(4, maxLocalOneRemote, Arrays.asList(dc1Dn1, dc1Dn2, dc1Dn3, dc2Dn2), 0);
    MockPartitionId majorDc21 =
        new MockPartitionId(5, maxLocalOneRemote, Arrays.asList(dc2Dn1, dc2Dn2, dc2Dn3, dc1Dn1), 0);
    MockPartitionId majorDc22 =
        new MockPartitionId(6, maxLocalOneRemote, Arrays.asList(dc2Dn1, dc2Dn2, dc2Dn3, dc1Dn2), 0);

    Collection<MockPartitionId> allPartitionIdsMain = Collections.unmodifiableSet(
        new HashSet<>(Arrays.asList(everywhere1, everywhere2, majorDc11, majorDc12, majorDc21, majorDc22)));
    ClusterManagerQueryHelper mockClusterManagerQueryHelper = Mockito.mock(ClusterManagerQueryHelper.class);
    doReturn(allPartitionIdsMain).when(mockClusterManagerQueryHelper).getPartitions();
    doReturn(false).when(mockClusterManagerQueryHelper).isPartitionFilteringEnabled();
    ClusterMapUtils.PartitionSelectionHelper psh =
        new ClusterMapUtils.PartitionSelectionHelper(mockClusterManagerQueryHelper, null, minimumLocalReplicaCount,
            maxReplicasAllSites, null);

    String[] dcsToTry = {null, "", dc1, dc2};
    for (String dc : dcsToTry) {
      Set<MockPartitionId> allPartitionIds = new HashSet<>(allPartitionIdsMain);
      resetPartitions(allPartitionIds);
      psh.updatePartitions(allPartitionIds, dc);

      // getPartitions()
      assertCollectionEquals("Partitions returned not as expected", allPartitionIds, psh.getPartitions(null));
      assertCollectionEquals("Partitions returned for " + maxReplicasAllSites + " not as expected",
          Arrays.asList(everywhere1, everywhere2), psh.getPartitions(maxReplicasAllSites));
      assertCollectionEquals("Partitions returned for " + maxLocalOneRemote + " not as expected",
          Arrays.asList(majorDc11, majorDc12, majorDc21, majorDc22), psh.getPartitions(maxLocalOneRemote));
      checkCaseInsensitivityForPartitionSelectionHelper(psh, true, maxReplicasAllSites,
          Arrays.asList(everywhere1, everywhere2));
      checkCaseInsensitivityForPartitionSelectionHelper(psh, true, maxLocalOneRemote,
          Arrays.asList(majorDc11, majorDc12, majorDc21, majorDc22));

      assertCollectionEquals("Default partition class should be used if partition class is unrecognized",
          Arrays.asList(everywhere1, everywhere2), psh.getPartitions(getRandomString(3)));

      // getWritablePartitions()
      Set<MockPartitionId> expectedWritableForMaxLocalOneRemote = null;
      MockPartitionId candidate1 = null;
      MockPartitionId candidate2 = null;
      if (dc != null) {
        switch (dc) {
          case dc1:
            candidate1 = majorDc11;
            candidate2 = majorDc12;
            expectedWritableForMaxLocalOneRemote = new HashSet<>(Arrays.asList(majorDc11, majorDc12));
            break;
          case dc2:
            candidate1 = majorDc21;
            candidate2 = majorDc22;
            expectedWritableForMaxLocalOneRemote = new HashSet<>(Arrays.asList(majorDc21, majorDc22));
            break;
        }
      }
      // invalid class should fall back to default partition class
      verifyGetWritablePartition(psh, allPartitionIds, getRandomString(3),
          new HashSet<>(Arrays.asList(everywhere1, everywhere2)), dc);
      verifyGetRandomWritablePartition(psh, allPartitionIds, getRandomString(3),
          new HashSet<>(Arrays.asList(everywhere1, everywhere2)), dc);

      verifyWritablePartitionsReturned(psh, allPartitionIds, maxReplicasAllSites, everywhere1, everywhere2,
          maxLocalOneRemote, expectedWritableForMaxLocalOneRemote, dc);
      if (candidate1 != null && candidate2 != null) {
        verifyWritablePartitionsReturned(psh, allPartitionIds, maxLocalOneRemote, candidate1, candidate2,
            maxReplicasAllSites, new HashSet<>(Arrays.asList(everywhere1, everywhere2)), dc);
      }
    }
    // additional test: ensure getRandomWritablePartition now honors replica state for PUT request
    psh = new ClusterMapUtils.PartitionSelectionHelper(mockClusterManagerQueryHelper, dc1, minimumLocalReplicaCount,
        maxReplicasAllSites, null);
    ReplicaId replicaId = everywhere1.getReplicaIds()
        .stream()
        .filter(r -> r.getDataNodeId().getDatacenterName().equals(dc1))
        .findFirst()
        .get();
    // make one of local replicas OFFLINE for everywhere1
    everywhere1.setReplicaState(replicaId, ReplicaState.OFFLINE);
    // remove some partitions that have 3 eligible replicas in DC1
    assertEquals("The partition selection helper should choose everywhere2", everywhere2,
        psh.getRandomWritablePartition(maxReplicasAllSites, null));
  }
Tests for all functions in {@link ClusterMapUtils.PartitionSelectionHelper}
ambry-clustermap/src/test/java/com/github/ambry/clustermap/ClusterMapUtilsTest.java