@Test
  public void limitMaxPartitionCountPerRequestTest() throws Exception {
    MockClusterMap clusterMap = new MockClusterMap();
    Pair<MockHost, MockHost> localAndRemoteHosts = getLocalAndRemoteHosts(clusterMap);
    MockHost localHost = localAndRemoteHosts.getFirst();
    MockHost remoteHost = localAndRemoteHosts.getSecond();

    List<PartitionId> partitionIds = clusterMap.getAllPartitionIds(null);
    for (PartitionId partitionId : partitionIds) {
      // add 5 messages into each partition and place it on remote host only
      addPutMessagesToReplicasOfPartition(partitionId, Collections.singletonList(remoteHost), 5);
    }
    StoreKeyFactory storeKeyFactory = Utils.getObj("com.github.ambry.commons.BlobIdFactory", clusterMap);
    MockStoreKeyConverterFactory mockStoreKeyConverterFactory = new MockStoreKeyConverterFactory(null, null);
    mockStoreKeyConverterFactory.setReturnInputIfAbsent(true);
    mockStoreKeyConverterFactory.setConversionMap(new HashMap<>());
    // we set batchSize to 10 in order to get all messages from one partition within single replication cycle
    int batchSize = 10;
    StoreKeyConverter storeKeyConverter = mockStoreKeyConverterFactory.getStoreKeyConverter();
    Transformer transformer = new ValidatingTransformer(storeKeyFactory, storeKeyConverter);
    // we set max partition count per request to 5, which forces thread to replicate replicas in two cycles. (Note that
    // number of partition to replicate is 10, it's 9 default partition class and one special class. They will be
    // replicated in two batches)
    ReplicationConfig initialReplicationConfig = replicationConfig;
    properties.setProperty("replication.max.partition.count.per.request", String.valueOf(5));
    replicationConfig = new ReplicationConfig(new VerifiableProperties(properties));
    CountDownLatch replicationCompleted = new CountDownLatch(partitionIds.size());
    AtomicReference<Exception> exception = new AtomicReference<>();
    Pair<Map<DataNodeId, List<RemoteReplicaInfo>>, ReplicaThread> replicasAndThread =
        getRemoteReplicasAndReplicaThread(batchSize, clusterMap, localHost, storeKeyConverter, transformer,
            (store, messageInfos) -> {
              try {
                replicationCompleted.countDown();
                // for each partition, replication should complete within single cycle (fetch once should suffice), so
                // we shut down local store once blobs are written. This can avoid unnecessary metadata requests sent to
                // remote host.
                store.shutdown();
              } catch (Exception e) {
                exception.set(e);
              }
            }, null, remoteHost);
    ReplicaThread replicaThread = replicasAndThread.getSecond();
    Thread thread = Utils.newThread(replicaThread, false);
    thread.start();
    assertTrue("Replication didn't complete within 10 secs", replicationCompleted.await(10, TimeUnit.SECONDS));
    // verify the # of replicas per metadata request is limited to 5 (note that there are 10 replicas to replicate, they
    // are split into to 2 small batches and get replicated in separate requests)
    assertEquals("There should be 2 metadata requests and each has 5 replicas to replicate", Arrays.asList(5, 5),
        remoteHost.replicaCountPerRequestTracker);
    // shutdown
    replicaThread.shutdown();
    if (exception.get() != null) {
      throw exception.get();
    }
    replicationConfig = initialReplicationConfig;
  }
Test that max partition count per request is honored in {@link ReplicaThread} if there are too many partitions to
replicate from the remote node.
@throws Exception
ambry-replication/src/test/java/com/github/ambry/replication/ReplicationTest.java