Examples with State org.apache.gobblin.configuration.State used on opensource projects

Example 96 with State

use of org.apache.gobblin.configuration.State in project incubator-gobblin by apache.

the class PartitionedWriterTest method test.

@Test
public void test() 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));
    Assert.assertEquals(builder.actions.size(), 2);
    TestPartitionAwareWriterBuilder.Action action = builder.actions.poll();
    Assert.assertEquals(action.getPartition(), "a");
    Assert.assertEquals(action.getType(), TestPartitionAwareWriterBuilder.Actions.BUILD);
    action = builder.actions.poll();
    Assert.assertEquals(action.getPartition(), "a");
    Assert.assertEquals(action.getType(), TestPartitionAwareWriterBuilder.Actions.WRITE);
    Assert.assertEquals(action.getTarget(), record1);
    Assert.assertTrue(writer.isSpeculativeAttemptSafe());
    String record2 = "123";
    writer.writeEnvelope(new RecordEnvelope(record2));
    Assert.assertEquals(builder.actions.size(), 2);
    action = builder.actions.poll();
    Assert.assertEquals(action.getPartition(), "1");
    Assert.assertEquals(action.getType(), TestPartitionAwareWriterBuilder.Actions.BUILD);
    Assert.assertFalse(writer.isSpeculativeAttemptSafe());
    action = builder.actions.poll();
    Assert.assertEquals(action.getPartition(), "1");
    Assert.assertEquals(action.getType(), TestPartitionAwareWriterBuilder.Actions.WRITE);
    Assert.assertEquals(action.getTarget(), record2);
    writer.writeEnvelope(new RecordEnvelope(record1));
    Assert.assertEquals(builder.actions.size(), 1);
    action = builder.actions.poll();
    Assert.assertEquals(action.getPartition(), "a");
    Assert.assertEquals(action.getType(), TestPartitionAwareWriterBuilder.Actions.WRITE);
    Assert.assertEquals(action.getTarget(), record1);
    Assert.assertEquals(writer.recordsWritten(), 3);
    Assert.assertEquals(writer.bytesWritten(), 3);
    Assert.assertFalse(writer.isSpeculativeAttemptSafe());
    writer.cleanup();
    Assert.assertEquals(builder.actions.size(), 2);
    action = builder.actions.poll();
    Assert.assertEquals(action.getType(), TestPartitionAwareWriterBuilder.Actions.CLEANUP);
    action = builder.actions.poll();
    Assert.assertEquals(action.getType(), TestPartitionAwareWriterBuilder.Actions.CLEANUP);
    writer.close();
    Assert.assertEquals(builder.actions.size(), 2);
    action = builder.actions.poll();
    Assert.assertEquals(action.getType(), TestPartitionAwareWriterBuilder.Actions.CLOSE);
    action = builder.actions.poll();
    Assert.assertEquals(action.getType(), TestPartitionAwareWriterBuilder.Actions.CLOSE);
    writer.commit();
    Assert.assertEquals(builder.actions.size(), 2);
    action = builder.actions.poll();
    Assert.assertEquals(action.getType(), TestPartitionAwareWriterBuilder.Actions.COMMIT);
    action = builder.actions.poll();
    Assert.assertEquals(action.getType(), TestPartitionAwareWriterBuilder.Actions.COMMIT);
}

Example 97 with State

use of org.apache.gobblin.configuration.State 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();
}

Example 98 with State

use of org.apache.gobblin.configuration.State in project incubator-gobblin by apache.

the class RetryWriterTest method retryTest.

public void retryTest() throws IOException {
    DataWriter<Void> writer = mock(DataWriter.class);
    doThrow(new RuntimeException()).when(writer).writeEnvelope(any(RecordEnvelope.class));
    DataWriterWrapperBuilder<Void> builder = new DataWriterWrapperBuilder<>(writer, new State());
    DataWriter<Void> retryWriter = builder.build();
    try {
        retryWriter.writeEnvelope(new RecordEnvelope<>(null));
        Assert.fail("Should have failed.");
    } catch (Exception e) {
    }
    verify(writer, times(5)).writeEnvelope(any(RecordEnvelope.class));
}

Example 99 with State

use of org.apache.gobblin.configuration.State in project incubator-gobblin by apache.

the class RetryWriterTest method retryTestSuccess.

public void retryTestSuccess() throws IOException {
    DataWriter<Void> writer = mock(DataWriter.class);
    DataWriterWrapperBuilder<Void> builder = new DataWriterWrapperBuilder<>(writer, new State());
    DataWriter<Void> retryWriter = builder.build();
    retryWriter.writeEnvelope(new RecordEnvelope<>(null));
    verify(writer, times(1)).writeEnvelope(any(RecordEnvelope.class));
}

Example 100 with State

use of org.apache.gobblin.configuration.State in project incubator-gobblin by apache.

the class DatePartitionedAvroFileExtractorTest method verifyWorkUnits.

private void verifyWorkUnits(List<WorkUnit> workunits, int expectedSize) throws DataRecordException, IOException {
    for (int i = 0; i < expectedSize; i++) {
        WorkUnit workUnit = ((MultiWorkUnit) workunits.get(i)).getWorkUnits().get(0);
        WorkUnitState wuState = new WorkUnitState(workunits.get(i), new State());
        wuState.setProp(ConfigurationKeys.SOURCE_FILEBASED_FS_URI, ConfigurationKeys.LOCAL_FS_URI);
        wuState.setProp(ConfigurationKeys.SOURCE_FILEBASED_FILES_TO_PULL, workUnit.getProp(ConfigurationKeys.SOURCE_FILEBASED_FILES_TO_PULL));
        try (DatePartitionedAvroFileExtractor extractor = new DatePartitionedAvroFileExtractor(wuState)) {
            GenericRecord record = extractor.readRecord(null);
            Assert.assertEquals(recordTimestamps[i], record.get(PARTITION_COLUMN_NAME));
            Assert.assertEquals(recordTimestamps[i], workUnit.getPropAsLong(ConfigurationKeys.WORK_UNIT_DATE_PARTITION_KEY));
        }
    }
}