private void mockReplicaStates(PartitionId partition, int count) {
    // Create mock replica list with the specified count
    List<AmbryReplica> replicas = new ArrayList<>();
    for (int i = 0; i < count; i++) {
      replicas.add(mock(AmbryReplica.class));
    }

    // Create state map for the relevant states
    Map<ReplicaState, List<AmbryReplica>> stateMap = new HashMap<>();
    if (count > 0) {
      stateMap.put(ReplicaState.LEADER, Collections.singletonList(replicas.get(0)));

      List<AmbryReplica> standbyReplicas = new ArrayList<>();
      for (int i = 1; i < count; i++) {
        standbyReplicas.add(replicas.get(i));
      }

      if (!standbyReplicas.isEmpty()) {
        stateMap.put(ReplicaState.STANDBY, standbyReplicas);
      }
    }

    // Mock getReplicaIdsByStates method to return our state map
    when(partition.getReplicaIdsByStates(ArgumentMatchers.anySet(), eq(datacenterName)))
        .thenReturn((Map) stateMap);
  }
Helper method to mock replica states for a partition
@param partition The partition to mock
@param count The number of active replicas to simulate
ambry-prioritization/src/test/java/com/github/ambry/replica/prioritization/ReplicationPrioritizationManagerTest.java