@Test
  public void testBufferNotCloseOnFail() throws Exception {
    outputBuffers =
        outputBuffers
            .withBuffer("out", new UnpartitionedPagePartitionFunction())
            .withNoMoreBufferIds();
    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.fail(new Exception("test"));
    assertEquals(taskExecution.getTaskInfo(false).getState(), TaskState.FAILED);

    // buffer will not be closed by fail 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));
    assertFalse(bufferResult.isBufferClosed());

    bufferResult =
        taskExecution.getResults(
            "out", 0, new DataSize(1, Unit.MEGABYTE), new Duration(500, TimeUnit.MILLISECONDS));
    assertFalse(bufferResult.isBufferClosed());
  }
  @Test
  public void testBufferCloseOnFinish() throws Exception {
    outputBuffers =
        outputBuffers
            .withBuffer("out", new UnpartitionedPagePartitionFunction())
            .withNoMoreBufferIds();
    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.addSources(
        ImmutableList.of(new TaskSource(tableScanNodeId, ImmutableSet.<ScheduledSplit>of(), true)));
    assertEquals(taskExecution.getTaskInfo(false).getState(), TaskState.FINISHED);

    // buffer will be closed by cancel event (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());
  }