Source:
public void testFreeFailedSlots() throws Exception { try (SlotPoolImpl slotPool = createAndSetUpSlotPool()) { final int parallelism = 5; final ArrayBlockingQueue<AllocationID> allocationIds = new ArrayBlockingQueue<>(parallelism); resourceManagerGateway.setRequestSlotConsumer( slotRequest -> allocationIds.offer(slotRequest.getAllocationId())); final Map<SlotRequestId, CompletableFuture<PhysicalSlot>> slotRequestFutures = new HashMap<>(parallelism); for (int i = 0; i < parallelism; i++) { final SlotRequestId slotRequestId = new SlotRequestId(); slotRequestFutures.put(slotRequestId, requestNewAllocatedSlot(slotPool, slotRequestId)); } final List<SlotOffer> slotOffers = new ArrayList<>(parallelism); for (int i = 0; i < parallelism; i++) { slotOffers.add(new SlotOffer(allocationIds.take(), i, ResourceProfile.ANY)); } slotPool.registerTaskManager(taskManagerLocation.getResourceID()); slotPool.offerSlots(taskManagerLocation, taskManagerGateway, slotOffers); FutureUtils.waitForAll(slotRequestFutures.values()).get(); final ArrayBlockingQueue<AllocationID> freedSlots = new ArrayBlockingQueue<>(1); taskManagerGateway.setFreeSlotFunction( (allocationID, throwable) -> { freedSlots.offer(allocationID); return CompletableFuture.completedFuture(Acknowledge.get()); }); final FlinkException failException = new FlinkException("Test fail exception"); for (int i = 0; i < parallelism - 1; i++) { final SlotOffer slotOffer = slotOffers.get(i); Optional<ResourceID> emptyTaskExecutorFuture = slotPool.failAllocation(slotOffer.getAllocationId(), failException); assertThat(emptyTaskExecutorFuture.isPresent(), is(false)); assertThat(freedSlots.take(), is(equalTo(slotOffer.getAllocationId()))); } final SlotOffer slotOffer = slotOffers.get(parallelism - 1); final Optional<ResourceID> emptyTaskExecutorFuture = slotPool.failAllocation( slotOffer.getAllocationId(), failException); assertTrue(emptyTaskExecutorFuture.isPresent()); assertThat(emptyTaskExecutorFuture.get(), is(equalTo(taskManagerLocation.getResourceID()))); assertThat(freedSlots.take(), is(equalTo(slotOffer.getAllocationId()))); } }

Target:
public void testFreeFailedSlots() throws Exception { try (SlotPoolImpl slotPool = createAndSetUpSlotPool(resourceManagerGateway)) { final int parallelism = 5; final ArrayBlockingQueue<AllocationID> allocationIds = new ArrayBlockingQueue<>(parallelism); resourceManagerGateway.setRequestSlotConsumer( slotRequest -> allocationIds.offer(slotRequest.getAllocationId())); final Map<SlotRequestId, CompletableFuture<PhysicalSlot>> slotRequestFutures = new HashMap<>(parallelism); for (int i = 0; i < parallelism; i++) { final SlotRequestId slotRequestId = new SlotRequestId(); slotRequestFutures.put(slotRequestId, requestNewAllocatedSlot(slotPool, slotRequestId)); } final List<SlotOffer> slotOffers = new ArrayList<>(parallelism); for (int i = 0; i < parallelism; i++) { slotOffers.add(new SlotOffer(allocationIds.take(), i, ResourceProfile.ANY)); } slotPool.registerTaskManager(taskManagerLocation.getResourceID()); slotPool.offerSlots(taskManagerLocation, taskManagerGateway, slotOffers); FutureUtils.waitForAll(slotRequestFutures.values()).get(); final ArrayBlockingQueue<AllocationID> freedSlots = new ArrayBlockingQueue<>(1); taskManagerGateway.setFreeSlotFunction( (allocationID, throwable) -> { freedSlots.offer(allocationID); return CompletableFuture.completedFuture(Acknowledge.get()); }); final FlinkException failException = new FlinkException("Test fail exception"); for (int i = 0; i < parallelism - 1; i++) { final SlotOffer slotOffer = slotOffers.get(i); Optional<ResourceID> emptyTaskExecutorFuture = slotPool.failAllocation(slotOffer.getAllocationId(), failException); assertThat(emptyTaskExecutorFuture.isPresent(), is(false)); assertThat(freedSlots.take(), is(equalTo(slotOffer.getAllocationId()))); } final SlotOffer slotOffer = slotOffers.get(parallelism - 1); final Optional<ResourceID> emptyTaskExecutorFuture = slotPool.failAllocation( slotOffer.getAllocationId(), failException); assertThat(emptyTaskExecutorFuture.get(), is(equalTo(taskManagerLocation.getResourceID()))); assertThat(freedSlots.take(), is(equalTo(slotOffer.getAllocationId()))); } }