@Override
  public Map<ReplicaState, List<ReplicaId>> getReplicaIdsByStates(Set<ReplicaState> states, String dcName) {
    Map<ReplicaState, List<ReplicaId>> replicasByState = new HashMap<>();
    updateReplicaAndStateIfNeeded();
    for (Map.Entry<ReplicaId, ReplicaState> entry : replicaAndState.entrySet()) {
      if (states.contains(entry.getValue())) {
        replicasByState.computeIfAbsent(entry.getValue(), k -> new ArrayList<>()).add(entry.getKey());
      }
    }
    return replicasByState;
  }
Mock partition id for unit tests
ambry-test-utils/src/main/java/com/github/ambry/clustermap/MockPartitionId.java