@Test
  public void testExpiredRequestAllocationOnAnyHost() throws Exception {
    MockClusterResourceManager spyManager = spy(new MockClusterResourceManager(callback, state));
    ContainerManager spyContainerManager =
        spy(new ContainerManager(containerPlacementMetadataStore, state, spyManager, true, false, mock(LocalityManager.class), faultDomainManager, config));
    spyAllocator = Mockito.spy(
        new ContainerAllocator(spyManager, config, state, true, spyContainerManager));
    // Request Preferred Resources
    spyAllocator.requestResources(new HashMap<String, String>() {
      {
        put("0", "hostname-0");
        put("1", "hostname-1");
      }
    });

    spyAllocatorThread = new Thread(spyAllocator);
    // Start the container allocator thread periodic assignment
    spyAllocatorThread.start();

    // Let the preferred host requests and the follow-up ANY_HOST request expire, expiration timeout is 500 ms
    Thread.sleep(1500);

    // Verify that all the request that were created as preferred host requests expired
    assertTrue(state.preferredHostRequests.get() == 2);
    assertTrue(state.expiredPreferredHostRequests.get() == 2);
    // expirations for initial preferred host requests
    verify(spyContainerManager).handleExpiredRequest(eq("0"), eq("hostname-0"),
        any(SamzaResourceRequest.class), any(ContainerAllocator.class), any(ResourceRequestState.class));
    verify(spyContainerManager).handleExpiredRequest(eq("1"), eq("hostname-1"),
        any(SamzaResourceRequest.class), any(ContainerAllocator.class), any(ResourceRequestState.class));
    // expirations for follow-up ANY_HOST requests
    // allocator keeps running in a loop so it might expire more than once during the wait time
    verify(spyContainerManager, atLeast(1)).handleExpiredRequest(eq("0"), eq(ResourceRequestState.ANY_HOST),
        any(SamzaResourceRequest.class), any(ContainerAllocator.class), any(ResourceRequestState.class));
    verify(spyContainerManager, atLeast(1)).handleExpiredRequest(eq("1"), eq(ResourceRequestState.ANY_HOST),
        any(SamzaResourceRequest.class), any(ContainerAllocator.class), any(ResourceRequestState.class));

    // Verify that preferred host request were cancelled and since no surplus resources were available
    // requestResource was invoked with ANY_HOST requests
    ArgumentCaptor<SamzaResourceRequest> cancelledRequestCaptor = ArgumentCaptor.forClass(SamzaResourceRequest.class);
    // preferred host requests and ANY_HOST requests got cancelled
    verify(spyManager, atLeast(4)).cancelResourceRequest(cancelledRequestCaptor.capture());
    assertEquals(ImmutableSet.of("hostname-0", "hostname-1", ResourceRequestState.ANY_HOST),
        cancelledRequestCaptor.getAllValues()
            .stream()
            .map(SamzaResourceRequest::getPreferredHost)
            .collect(Collectors.toSet()));
    assertTrue(state.matchedResourceRequests.get() == 0);
    // at least 2 for expired preferred host requests, 2 for expired follow-up ANY_HOST requests
    assertTrue(state.anyHostRequests.get() >= 4);
    spyAllocator.stop();
  }
Handles expired requests correctly and assigns ANY_HOST
samza-core/src/test/java/org/apache/samza/clustermanager/TestContainerAllocatorWithHostAffinity.java