HashMap<Object, Container> taskAllocations = new LinkedHashMap<Object, Container>();
PriorityBlockingQueue<TaskRequest> taskRequestQueue = new PriorityBlockingQueue<TaskRequest>();
TaskSchedulerAppCallback appClientDelegate = mock(TaskSchedulerAppCallback.class);
// Object under test
AsyncDelegateRequestHandler requestHandler =
new AsyncDelegateRequestHandler(taskRequestQueue,
containerFactory,
taskAllocations,
appClientDelegate,
tezConf);
// Allocate up to max tasks
for (int i = 0; i < MAX_TASKS; i++) {
Priority priority = Priority.newInstance(20);
requestHandler.addAllocateTaskRequest(new Long(i), null, priority, null);
requestHandler.processRequest();
}
// Only MAX_TASKS number of tasks should have been allocated
Assert.assertEquals("Wrong number of allocate tasks", MAX_TASKS, taskAllocations.size());
Assert.assertTrue("Another allocation should not fit", requestHandler.shouldWait());
// Deallocate down to zero
for (int i = 0; i < MAX_TASKS; i++) {
requestHandler.addDeallocateTaskRequest(new Long(i));
requestHandler.processRequest();
}
// All allocated tasks should have been removed
Assert.assertEquals("Wrong number of allocate tasks", 0, taskAllocations.size());
}