private void blockRebalanceConfirmation(
      final EmbeddedCacheManager manager, final CheckPoint checkPoint) throws Exception {
    ClusterTopologyManager ctm =
        TestingUtil.extractGlobalComponent(manager, ClusterTopologyManager.class);
    final Answer<Object> forwardedAnswer = AdditionalAnswers.delegatesTo(ctm);
    ClusterTopologyManager mockManager =
        mock(ClusterTopologyManager.class, withSettings().defaultAnswer(forwardedAnswer));
    doAnswer(
            new Answer<Object>() {
              @Override
              public Object answer(InvocationOnMock invocation) throws Throwable {
                Object[] arguments = invocation.getArguments();
                Address source = (Address) arguments[1];
                int topologyId = (Integer) arguments[2];
                checkPoint.trigger("pre_rebalance_confirmation_" + topologyId + "_from_" + source);
                checkPoint.awaitStrict(
                    "resume_rebalance_confirmation_" + topologyId + "_from_" + source, 10, SECONDS);

                return forwardedAnswer.answer(invocation);
              }
            })
        .when(mockManager)
        .handleRebalanceCompleted(
            anyString(), any(Address.class), anyInt(), any(Throwable.class), anyInt());
    TestingUtil.replaceComponent(manager, ClusterTopologyManager.class, mockManager, true);
  }
 private void blockJoinResponse(final EmbeddedCacheManager manager, final CheckPoint checkPoint)
     throws Exception {
   ClusterTopologyManager ctm =
       TestingUtil.extractGlobalComponent(manager, ClusterTopologyManager.class);
   final Answer<Object> forwardedAnswer = AdditionalAnswers.delegatesTo(ctm);
   ClusterTopologyManager mockManager =
       mock(ClusterTopologyManager.class, withSettings().defaultAnswer(forwardedAnswer));
   doAnswer(
           new Answer<Object>() {
             @Override
             public Object answer(InvocationOnMock invocation) throws Throwable {
               Object answer = forwardedAnswer.answer(invocation);
               checkPoint.trigger("sending_join_response");
               checkPoint.awaitStrict("resume_join_response", 10, SECONDS);
               return answer;
             }
           })
       .when(mockManager)
       .handleJoin(anyString(), any(Address.class), any(CacheJoinInfo.class), anyInt());
   TestingUtil.replaceComponent(manager, ClusterTopologyManager.class, mockManager, true);
 }
  protected void waitUntilStateBeingTransferred(
      final Cache<?, ?> cache, final CheckPoint checkPoint) {
    StateConsumer sc = TestingUtil.extractComponent(cache, StateConsumer.class);
    final Answer<Object> forwardedAnswer = AdditionalAnswers.delegatesTo(sc);
    StateConsumer mockConsumer =
        mock(StateConsumer.class, withSettings().defaultAnswer(forwardedAnswer));
    doAnswer(
            new Answer() {
              @Override
              public Object answer(InvocationOnMock invocation) throws Throwable {
                // Wait for main thread to sync up
                checkPoint.trigger("pre_state_apply_invoked_for_" + cache);
                // Now wait until main thread lets us through
                checkPoint.awaitStrict(
                    "pre_state_apply_release_for_" + cache, 10, TimeUnit.SECONDS);

                return forwardedAnswer.answer(invocation);
              }
            })
        .when(mockConsumer)
        .applyState(any(Address.class), anyInt(), anyCollection());
    TestingUtil.replaceComponent(cache, StateConsumer.class, mockConsumer, true);
  }
package org.infinispan.iteration;