Source:
public void testOnInitToForceRestartAMHighAvailability() throws Exception { Map<String, String> configMap = new HashMap<>(configVals); configMap.put(JobConfig.YARN_AM_HIGH_AVAILABILITY_ENABLED, "true"); Config conf = new MapConfig(configMap); SamzaApplicationState state = new SamzaApplicationState(getJobModelManager(2)); state.runningProcessors.put("0", new SamzaResource(1, 1024, "host", "0")); MockClusterResourceManagerCallback callback = new MockClusterResourceManagerCallback(); ClusterResourceManager clusterResourceManager = new MockClusterResourceManager(callback, state); ClusterManagerConfig clusterManagerConfig = spy(new ClusterManagerConfig(conf)); ContainerManager containerManager = buildContainerManager(containerPlacementMetadataStore, state, clusterResourceManager, clusterManagerConfig.getHostAffinityEnabled(), false); ContainerProcessManager cpm = buildContainerProcessManager(clusterManagerConfig, state, clusterResourceManager, Optional.empty(), true); MockContainerAllocatorWithoutHostAffinity allocator = new MockContainerAllocatorWithoutHostAffinity( clusterResourceManager, conf, state, containerManager); getPrivateFieldFromCpm("containerAllocator", cpm).set(cpm, allocator); CountDownLatch latch = new CountDownLatch(1); getPrivateFieldFromCpm("allocatorThread", cpm).set(cpm, new Thread() { public void run() { isRunning = true; latch.countDown(); } }); cpm.start(); if (!latch.await(2, TimeUnit.SECONDS)) { Assert.fail("timed out waiting for the latch to expire"); } assertEquals("CPM should stop the running container", 1, callback.resourceStatuses.size()); SamzaResourceStatus actualResourceStatus = callback.resourceStatuses.get(0); assertEquals("Container 0 should be stopped", "0", actualResourceStatus.getContainerId()); assertEquals("Container 0 should have exited with preempted status", SamzaResourceStatus.PREEMPTED, actualResourceStatus.getExitCode()); cpm.stop(); }

Target:
public void testOnInitToForceRestartAMHighAvailability() throws Exception { Map<String, String> configMap = new HashMap<>(configVals); configMap.put(JobConfig.YARN_AM_HIGH_AVAILABILITY_ENABLED, "true"); Config conf = new MapConfig(configMap); SamzaResource samzaResource = new SamzaResource(1, 1024, "host", "0"); SamzaApplicationState state = new SamzaApplicationState(getJobModelManager(2)); state.runningProcessors.put("0", samzaResource); MockClusterResourceManagerCallback callback = new MockClusterResourceManagerCallback(); ClusterResourceManager clusterResourceManager = spy(new MockClusterResourceManager(callback, state)); ClusterManagerConfig clusterManagerConfig = spy(new ClusterManagerConfig(conf)); ContainerManager containerManager = buildContainerManager(containerPlacementMetadataStore, state, clusterResourceManager, clusterManagerConfig.getHostAffinityEnabled(), false); ContainerProcessManager cpm = buildContainerProcessManager(clusterManagerConfig, state, clusterResourceManager, Optional.empty(), true); MockContainerAllocatorWithoutHostAffinity allocator = new MockContainerAllocatorWithoutHostAffinity( clusterResourceManager, conf, state, containerManager); getPrivateFieldFromCpm("containerAllocator", cpm).set(cpm, allocator); CountDownLatch latch = new CountDownLatch(1); getPrivateFieldFromCpm("allocatorThread", cpm).set(cpm, new Thread() { public void run() { isRunning = true; latch.countDown(); } }); cpm.start(); if (!latch.await(2, TimeUnit.SECONDS)) { Assert.fail("timed out waiting for the latch to expire"); } verify(clusterResourceManager, times(1)).stopStreamProcessor(samzaResource); assertEquals("CPM should stop the running container", 1, callback.resourceStatuses.size()); SamzaResourceStatus actualResourceStatus = callback.resourceStatuses.get(0); assertEquals("Container 0 should be stopped", "0", actualResourceStatus.getContainerId()); assertEquals("Container 0 should have exited with preempted status", SamzaResourceStatus.PREEMPTED, actualResourceStatus.getExitCode()); cpm.stop(); }