@Test
  public void testReestablishConnectionWithNewAM() throws InterruptedException {
    String newCoordinatorUrl = "http://some-host-2.prod.linkedin.com";
    this.containerHeartbeatMonitor = spy(buildContainerHeartbeatMonitor(true));
    CoordinatorStreamValueSerde serde = new CoordinatorStreamValueSerde(SetConfig.TYPE);
    ContainerHeartbeatMonitor.ContainerHeartbeatMetrics metrics = this.containerHeartbeatMonitor.getMetrics();

    when(this.containerHeartbeatClient.requestHeartbeat()).thenReturn(FAILURE_RESPONSE).thenReturn(SUCCESS_RESPONSE);
    when(this.containerHeartbeatMonitor
        .createContainerHeartbeatClient(newCoordinatorUrl, CONTAINER_EXECUTION_ID))
        .thenReturn(this.containerHeartbeatClient);
    when(this.coordinatorStreamStore.get(CoordinationConstants.YARN_COORDINATOR_URL))
        .thenReturn(serde.toBytes(newCoordinatorUrl));

    this.containerHeartbeatMonitor.start();
    // wait for the executor to finish the heartbeat check task
    boolean fixedRateTaskCompleted = this.schedulerFixedRateExecutionLatch.await(2, TimeUnit.SECONDS);
    assertTrue("Did not complete heartbeat check", fixedRateTaskCompleted);
    assertEquals("Heartbeat expired count should be 1", 1, metrics.getHeartbeatExpiredCount().getCount());
    assertEquals("Heartbeat established failure count should be 0", 0,
        metrics.getHeartbeatEstablishedFailureCount().getCount());
    assertEquals("Heartbeat established with new AM should be 1", 1,
        metrics.getHeartbeatEstablishedWithNewAmCount().getCount());
    // shutdown task should not have been submitted
    verify(this.scheduler, never()).schedule(any(Runnable.class), anyLong(), any());
    verify(this.onExpired, never()).run();

    this.containerHeartbeatMonitor.stop();
    verify(this.scheduler).shutdown();
  }
Use this to detect when the scheduler has finished executing the fixed-rate task.
samza-core/src/test/java/org/apache/samza/container/TestContainerHeartbeatMonitor.java