@Test
  public void testRowNumWithThrottling() throws Exception {
    ArgumentCaptor<AggregateResponse> responseCaptor =
      ArgumentCaptor.forClass(AggregateResponse.class);
    RpcCallback<AggregateResponse> callback = mock(RpcCallback.class);

    // First call should get throttled
    aggregate.getRowNum(controller, request, callback);

    verify(callback).run(responseCaptor.capture());

    AggregateResponse response = responseCaptor.getValue();
    assertTrue("Response should indicate there are more rows", response.hasNextChunkStartRow());
    assertEquals("Wait interval should be set", 1000, response.getWaitIntervalMs());
    assertEquals(THROTTLE_AT_ROW - 1, response.getFirstPart(0).asReadOnlyByteBuffer().getLong());

    AggregateRequest request2 = AggregateRequest.newBuilder(request)
      .setScan(request.getScan().toBuilder().setStartRow(response.getNextChunkStartRow()).build())
      .build();

    // Reset throttle for second call
    reset();

    RpcCallback<AggregateResponse> callback2 = mock(RpcCallback.class);
    aggregate.getRowNum(controller, request2, callback2);

    verify(callback2).run(responseCaptor.capture());

    AggregateResponse response2 = responseCaptor.getValue();
    assertEquals("Final row count should be correct", NUM_ROWS - THROTTLE_AT_ROW + 1,
      response2.getFirstPart(0).asReadOnlyByteBuffer().getLong());
    assertFalse("Response should not indicate there are more rows",
      response2.hasNextChunkStartRow());
  }
Test AggregateImplementation with throttling and partial results
hbase-endpoint/src/test/java/org/apache/hadoop/hbase/coprocessor/TestAggregateImplementation.java