Examples with BasicAckableForTesting org.apache.gobblin.ack.BasicAckableForTesting used on opensource projects

Example 6 with BasicAckableForTesting

use of org.apache.gobblin.ack.BasicAckableForTesting in project incubator-gobblin by apache.

the class TestRecordStream method testAcks.

@Test
public void testAcks() throws Exception {
    StreamEntity[] entities = new StreamEntity[] { new RecordEnvelope<>("a"), new BasicTestControlMessage("1"), new RecordEnvelope<>("b"), new BasicTestControlMessage("2") };
    BasicAckableForTesting ackable = new BasicAckableForTesting();
    for (int i = 0; i < entities.length; i++) {
        entities[i].addCallBack(ackable);
    }
    MyExtractor extractor = new MyExtractor(entities);
    MyConverter converter = new MyConverter();
    MyDataWriter writer = new MyDataWriter();
    // Create a TaskState
    TaskState taskState = getEmptyTestTaskState("testRetryTaskId");
    taskState.setProp(ConfigurationKeys.TASK_SYNCHRONOUS_EXECUTION_MODEL_KEY, false);
    // Create a mock TaskContext
    TaskContext mockTaskContext = mock(TaskContext.class);
    when(mockTaskContext.getExtractor()).thenReturn(extractor);
    when(mockTaskContext.getForkOperator()).thenReturn(new IdentityForkOperator());
    when(mockTaskContext.getTaskState()).thenReturn(taskState);
    when(mockTaskContext.getConverters()).thenReturn(Lists.newArrayList(converter));
    when(mockTaskContext.getTaskLevelPolicyChecker(any(TaskState.class), anyInt())).thenReturn(mock(TaskLevelPolicyChecker.class));
    when(mockTaskContext.getRowLevelPolicyChecker()).thenReturn(new RowLevelPolicyChecker(Lists.newArrayList(), "ss", FileSystem.getLocal(new Configuration())));
    when(mockTaskContext.getRowLevelPolicyChecker(anyInt())).thenReturn(new RowLevelPolicyChecker(Lists.newArrayList(), "ss", FileSystem.getLocal(new Configuration())));
    when(mockTaskContext.getDataWriterBuilder(anyInt(), anyInt())).thenReturn(writer);
    // Create a mock TaskPublisher
    TaskPublisher mockTaskPublisher = mock(TaskPublisher.class);
    when(mockTaskPublisher.canPublish()).thenReturn(TaskPublisher.PublisherState.SUCCESS);
    when(mockTaskContext.getTaskPublisher(any(TaskState.class), any(TaskLevelPolicyCheckResults.class))).thenReturn(mockTaskPublisher);
    // Create a mock TaskStateTracker
    TaskStateTracker mockTaskStateTracker = mock(TaskStateTracker.class);
    // Create a TaskExecutor - a real TaskExecutor must be created so a Fork is run in a separate thread
    TaskExecutor taskExecutor = new TaskExecutor(new Properties());
    // Create the Task
    Task realTask = new Task(mockTaskContext, mockTaskStateTracker, taskExecutor, Optional.<CountDownLatch>absent());
    Task task = spy(realTask);
    doNothing().when(task).submitTaskCommittedEvent();
    task.run();
    task.commit();
    Assert.assertEquals(task.getTaskState().getWorkingState(), WorkUnitState.WorkingState.SUCCESSFUL);
    Assert.assertEquals(ackable.acked, 4);
}

Example 7 with BasicAckableForTesting

use of org.apache.gobblin.ack.BasicAckableForTesting in project incubator-gobblin by apache.

the class ConverterTest method testSingleOutputIterable.

@Test
public void testSingleOutputIterable() throws Exception {
    MyConverter converter = new MyConverter();
    BasicAckableForTesting ackable = new BasicAckableForTesting();
    RecordStreamWithMetadata<Integer, String> stream = new RecordStreamWithMetadata<>(Flowable.just(new RecordEnvelope<>(1)), GlobalMetadata.<String>builder().schema("schema").build()).mapRecords(r -> {
        r.addCallBack(ackable);
        return r;
    });
    List<StreamEntity<Integer>> outputRecords = Lists.newArrayList();
    converter.processStream(stream, new WorkUnitState()).getRecordStream().subscribe(outputRecords::add);
    Assert.assertEquals(outputRecords.size(), 1);
    // output record has not been acked
    Assert.assertEquals(ackable.acked, 0);
    outputRecords.get(0).ack();
    // output record acked
    Assert.assertEquals(ackable.acked, 1);
}

Example 8 with BasicAckableForTesting

use of org.apache.gobblin.ack.BasicAckableForTesting in project incubator-gobblin by apache.

the class ConverterTest method testEmptyOutputIterable.

@Test
public void testEmptyOutputIterable() throws Exception {
    MyConverter converter = new MyConverter();
    BasicAckableForTesting ackable = new BasicAckableForTesting();
    RecordStreamWithMetadata<Integer, String> stream = new RecordStreamWithMetadata<>(Flowable.just(new RecordEnvelope<>(0)), GlobalMetadata.<String>builder().schema("schema").build()).mapRecords(r -> {
        r.addCallBack(ackable);
        return r;
    });
    List<StreamEntity<Integer>> outputRecords = Lists.newArrayList();
    converter.processStream(stream, new WorkUnitState()).getRecordStream().subscribe(outputRecords::add);
    Assert.assertEquals(outputRecords.size(), 0);
    // record got filtered, acked immediately
    Assert.assertEquals(ackable.acked, 1);
}

Example 9 with BasicAckableForTesting

use of org.apache.gobblin.ack.BasicAckableForTesting in project incubator-gobblin by apache.

the class RecordEnvelopeTest method testMultipleDerivedRecords.

@Test
public void testMultipleDerivedRecords() {
    BasicAckableForTesting ackable = new BasicAckableForTesting();
    RecordEnvelope<String> record = new RecordEnvelope<>("test", new MyWatermark(105));
    record.addCallBack(ackable);
    RecordEnvelope<String>.ForkRecordBuilder<String> forkRecordBuilder = record.forkRecordBuilder();
    RecordEnvelope<String> derivedRecord = forkRecordBuilder.childRecord("testDerived");
    derivedRecord.ack();
    // not acked yet as forkRecordBuilder has not been closed
    Assert.assertEquals(ackable.acked, 0);
    Assert.assertEquals(((MyWatermark) derivedRecord.getWatermark()).getId(), 105);
    RecordEnvelope<String> derivedRecord2 = forkRecordBuilder.childRecord("testDerived2");
    derivedRecord2.ack();
    forkRecordBuilder.close();
    Assert.assertEquals(ackable.acked, 1);
    Assert.assertEquals(((MyWatermark) derivedRecord2.getWatermark()).getId(), 105);
}

Example 10 with BasicAckableForTesting

use of org.apache.gobblin.ack.BasicAckableForTesting in project incubator-gobblin by apache.

the class RecordEnvelopeTest method testDerivedRecordCreation.

@Test
public void testDerivedRecordCreation() {
    BasicAckableForTesting ackable = new BasicAckableForTesting();
    RecordEnvelope<String> record = new RecordEnvelope<>("test", new MyWatermark(101));
    record.addCallBack(ackable);
    RecordEnvelope<String> derivedRecord = record.withRecord("testDerived");
    derivedRecord.ack();
    Assert.assertEquals(ackable.acked, 1);
    Assert.assertEquals(((MyWatermark) derivedRecord.getWatermark()).getId(), 101);
}