Source:
<code>public void testTaskManagerUnregistration() throws Exception { final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder() .setRequestSlotFunction(tuple6 -> new CompletableFuture<>()) .createTestingTaskExecutorGateway(); final TaskExecutorConnection taskManagerConnection = createTaskExecutorConnection(taskExecutorGateway); final ResourceID resourceId = taskManagerConnection.getResourceID(); final SlotID slotId1 = new SlotID(resourceId, 0); final SlotID slotId2 = new SlotID(resourceId, 1); final SlotReport slotReport = new SlotReport(Arrays.asList(createAllocatedSlotStatus(slotId1), createFreeSlotStatus(slotId2))); final ResourceRequirements resourceRequirements = createResourceRequirementsForSingleSlot(); final DefaultSlotTracker slotTracker = new DefaultSlotTracker(); try (DeclarativeSlotManager slotManager = createDeclarativeSlotManagerBuilder() .setSlotTracker(slotTracker) .buildAndStartWithDirectExec()) { slotManager.registerTaskManager(taskManagerConnection, slotReport); assertEquals("The number registered slots does not equal the expected number.", 2, slotManager.getNumberRegisteredSlots()); <START> DeclarativeTaskManagerSlot slot1 = slotTracker.getSlot(slotId1); DeclarativeTaskManagerSlot slot2 = slotTracker.getSlot(slotId2); assertSame(SlotState.ALLOCATED, slot1.getState()); assertSame(SlotState.FREE, slot2.getState()); <END> slotManager.processResourceRequirements(resourceRequirements); assertSame(SlotState.PENDING, slot2.getState()); slotManager.unregisterTaskManager(taskManagerConnection.getInstanceID(), TEST_EXCEPTION); assertEquals(0, slotManager.getNumberRegisteredSlots()); } }</code><technical_language>Do internal state of slots in order ensure a TM unregistered more slots? simply slotManager.getNumberRegisteredSlots() for this</technical_language>

Target:
public void testTaskManagerUnregistration() throws Exception { final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder() .setRequestSlotFunction(tuple6 -> new CompletableFuture<>()) .createTestingTaskExecutorGateway(); final TaskExecutorConnection taskManagerConnection = createTaskExecutorConnection(taskExecutorGateway); final ResourceID resourceId = taskManagerConnection.getResourceID(); final SlotID slotId1 = new SlotID(resourceId, 0); final SlotID slotId2 = new SlotID(resourceId, 1); final SlotReport slotReport = new SlotReport(Arrays.asList(createAllocatedSlotStatus(slotId1), createFreeSlotStatus(slotId2))); final ResourceRequirements resourceRequirements = createResourceRequirementsForSingleSlot(); final DefaultSlotTracker slotTracker = new DefaultSlotTracker(); try (DeclarativeSlotManager slotManager = createDeclarativeSlotManagerBuilder() .setSlotTracker(slotTracker) .buildAndStartWithDirectExec()) { slotManager.registerTaskManager(taskManagerConnection, slotReport); assertEquals("The number registered slots does not equal the expected number.", 2, slotManager.getNumberRegisteredSlots()); slotManager.processResourceRequirements(resourceRequirements); slotManager.unregisterTaskManager(taskManagerConnection.getInstanceID(), TEST_EXCEPTION); assertEquals(0, slotManager.getNumberRegisteredSlots()); } }