@Test
  public void testRequestAllocationOnPreferredHostWithRunStreamProcessor() throws Exception {
    ClusterResourceManager.Callback mockCPM = mock(MockClusterResourceManagerCallback.class);
    ClusterResourceManager mockClusterResourceManager = new MockClusterResourceManager(mockCPM, state);
    ContainerManager containerManager =
        new ContainerManager(containerPlacementMetadataStore, state, mockClusterResourceManager, true, false, mock(LocalityManager.class), faultDomainManager, config);
    // Mock the callback from ClusterManager to add resources to the allocator
    doAnswer((InvocationOnMock invocation) -> {
      SamzaResource resource = (SamzaResource) invocation.getArgumentAt(0, List.class).get(0);
      spyAllocator.addResource(resource);
      return null;
    }).when(mockCPM).onResourcesAvailable(anyList());

    spyAllocator = Mockito.spy(
        new ContainerAllocator(mockClusterResourceManager, config, state, true, containerManager));

    // Request Resources
    spyAllocator.requestResources(new HashMap<String, String>() {
      {
        put("0", "abc");
        put("1", "xyz");
      }
    });

    spyAllocatorThread = new Thread(spyAllocator);

    // Start the container allocator thread periodic assignment
    spyAllocatorThread.start();
    // Let Allocator thread periodically fulfill requests
    Thread.sleep(100);

    // Verify that all the request that were created were preferred host requests
    ArgumentCaptor<SamzaResourceRequest> resourceRequestCaptor = ArgumentCaptor.forClass(SamzaResourceRequest.class);
    verify(spyAllocator, times(2)).runStreamProcessor(resourceRequestCaptor.capture(), anyString());
    resourceRequestCaptor.getAllValues()
        .forEach(resourceRequest -> assertNotEquals(resourceRequest.getPreferredHost(), ResourceRequestState.ANY_HOST));
    Set<String> hostNames = resourceRequestCaptor.getAllValues().stream().map(request -> request.getPreferredHost()).collect(
        Collectors.toSet());
    assertTrue(hostNames.contains("abc"));
    assertTrue(hostNames.contains("xyz"));
    // No any host requests should be made if preferred host is satisfied
    assertTrue(state.anyHostRequests.get() == 0);
    // State check when host affinity is enabled
    assertTrue(state.matchedResourceRequests.get() == 2);
    assertTrue(state.preferredHostRequests.get() == 2);
    spyAllocator.stop();
  }
Handles expired requests correctly and assigns ANY_HOST
samza-core/src/test/java/org/apache/samza/clustermanager/TestContainerAllocatorWithHostAffinity.java