@Test
public void testCloseAfterFlush()
throws IOException {
WorkUnitState state = new WorkUnitState();
state.getJobState().setProp(CloseOnFlushWriterWrapper.WRITER_CLOSE_ON_FLUSH_KEY, "true");
List<DummyWriter> dummyWriters = new ArrayList<>();
CloseOnFlushWriterWrapper<byte[]> writer = getCloseOnFlushWriter(dummyWriters, state);
byte[] record = new byte[]{'a', 'b', 'c', 'd'};
writer.writeEnvelope(new RecordEnvelope(record));
writer.getMessageHandler().handleMessage(FlushControlMessage.builder().build());
Assert.assertEquals(dummyWriters.get(0).recordsWritten(), 1);
Assert.assertEquals(dummyWriters.get(0).flushCount, 1);
Assert.assertEquals(dummyWriters.get(0).closeCount, 1);
Assert.assertTrue(dummyWriters.get(0).committed);
Assert.assertEquals(dummyWriters.get(0).handlerCalled, 1);
writer.close();
Assert.assertEquals(dummyWriters.get(0).closeCount, 1);
}