Examples with FlushControlMessage org.apache.gobblin.stream.FlushControlMessage used on opensource projects

Example 1 with FlushControlMessage

use of org.apache.gobblin.stream.FlushControlMessage in project incubator-gobblin by apache.

the class PartitionedWriterTest method testControlMessageHandler.

@Test
public void testControlMessageHandler() throws IOException {
    State state = new State();
    state.setProp(ConfigurationKeys.WRITER_PARTITIONER_CLASS, TestPartitioner.class.getCanonicalName());
    TestPartitionAwareWriterBuilder builder = new TestPartitionAwareWriterBuilder();
    PartitionedDataWriter writer = new PartitionedDataWriter<String, String>(builder, state);
    Assert.assertEquals(builder.actions.size(), 0);
    String record1 = "abc";
    writer.writeEnvelope(new RecordEnvelope(record1));
    String record2 = "123";
    writer.writeEnvelope(new RecordEnvelope(record2));
    FlushControlMessage controlMessage = FlushControlMessage.builder().build();
    BasicAckableForTesting ackable = new BasicAckableForTesting();
    controlMessage.addCallBack(ackable);
    Assert.assertEquals(ackable.acked, 0);
    // when the control message is cloned properly then this does not raise an error
    writer.getMessageHandler().handleMessage(controlMessage);
    // message handler does not ack since consumeRecordStream does acking for control messages
    // this should be revisited when control message error handling is changed
    controlMessage.ack();
    Assert.assertEquals(ackable.acked, 1);
    writer.close();
}