@Test
  public void testCpmShouldStopWhenContainersFinish() throws Exception {
    Config conf = getConfig();
    SamzaApplicationState state = new SamzaApplicationState(getJobModelManager(1));
    MockClusterResourceManagerCallback callback = new MockClusterResourceManagerCallback();
    MockClusterResourceManager clusterResourceManager = new MockClusterResourceManager(callback, state);
    ClusterManagerConfig clusterManagerConfig = spy(new ClusterManagerConfig(conf));
    ContainerManager containerManager =
        buildContainerManager(containerPlacementMetadataStore, state, clusterResourceManager,
            clusterManagerConfig.getHostAffinityEnabled(), false);

    MockContainerAllocatorWithoutHostAffinity allocator = new MockContainerAllocatorWithoutHostAffinity(
        clusterResourceManager,
        conf,
        state,
        containerManager);

    ContainerProcessManager cpm =
        spy(buildContainerProcessManager(clusterManagerConfig, state, clusterResourceManager, Optional.of(allocator)));

    // start triggers a request
    cpm.start();

    assertFalse(cpm.shouldShutdown());
    assertEquals(1, allocator.getContainerRequestState().numPendingRequests());
    assertEquals(0, allocator.getContainerRequestState().numDelayedRequests());

    SamzaResource container = new SamzaResource(1, 1024, "host1", "id0");
    cpm.onResourceAllocated(container);

    // Allow container to run and update state

    if (!allocator.awaitContainersStart(1, 2, TimeUnit.SECONDS)) {
      fail("timed out waiting for the containers to start");
    }
    cpm.onStreamProcessorLaunchSuccess(container);
    assertFalse(cpm.shouldShutdown());

    cpm.onResourceCompleted(new SamzaResourceStatus("id0", "diagnostics", SamzaResourceStatus.SUCCESS));
    verify(cpm, never()).onResourceCompletedWithUnknownStatus(any(SamzaResourceStatus.class), anyString(), anyString(), anyInt());
    assertTrue(cpm.shouldShutdown());
    cpm.stop();
  }
Test Container Process Manager should stop when all containers finish
samza-core/src/test/java/org/apache/samza/clustermanager/TestContainerProcessManager.java