@Test public void testBufferCloseOnCancel() throws Exception { outputBuffers = outputBuffers.withBuffer("out", new UnpartitionedPagePartitionFunction()); taskExecution.addResultQueue(outputBuffers); BufferResult bufferResult = taskExecution.getResults( "out", 0, new DataSize(1, Unit.MEGABYTE), new Duration(0, TimeUnit.MILLISECONDS)); assertFalse(bufferResult.isBufferClosed()); bufferResult = taskExecution.getResults( "out", 0, new DataSize(1, Unit.MEGABYTE), new Duration(0, TimeUnit.MILLISECONDS)); assertFalse(bufferResult.isBufferClosed()); taskExecution.cancel(); assertEquals(taskExecution.getTaskInfo(false).getState(), TaskState.CANCELED); // buffer will be closed by cancel event. event is async so wait for 500 MS for event to fire bufferResult = taskExecution.getResults( "out", 0, new DataSize(1, Unit.MEGABYTE), new Duration(500, TimeUnit.MILLISECONDS)); assertTrue(bufferResult.isBufferClosed()); bufferResult = taskExecution.getResults( "out", 0, new DataSize(1, Unit.MEGABYTE), new Duration(500, TimeUnit.MILLISECONDS)); assertTrue(bufferResult.isBufferClosed()); }