Examples with BasicTestControlMessage org.apache.gobblin.runtime.BasicTestControlMessage used on opensource projects

Example 1 with BasicTestControlMessage

use of org.apache.gobblin.runtime.BasicTestControlMessage in project incubator-gobblin by apache.

the class ForkerTest method test.

@Test
public void test() throws Exception {
    Forker forker = new Forker();
    MyFlowable<StreamEntity<byte[]>> flowable = new MyFlowable<>();
    RecordStreamWithMetadata<byte[], String> stream = new RecordStreamWithMetadata<>(flowable, GlobalMetadata.<String>builder().schema("schema").build());
    WorkUnitState workUnitState = new WorkUnitState();
    workUnitState.setProp(ConfigurationKeys.FORK_BRANCHES_KEY, "3");
    Forker.ForkedStream<byte[], String> forkedStream = forker.forkStream(stream, new MyForkOperator(), workUnitState);
    Assert.assertEquals(forkedStream.getForkedStreams().size(), 3);
    Queue<StreamEntity<byte[]>> output0 = new LinkedList<>();
    forkedStream.getForkedStreams().get(0).getRecordStream().subscribe(output0::add);
    Queue<StreamEntity<byte[]>> output1 = new LinkedList<>();
    forkedStream.getForkedStreams().get(1).getRecordStream().subscribe(output1::add);
    Queue<StreamEntity<byte[]>> output2 = new LinkedList<>();
    forkedStream.getForkedStreams().get(2).getRecordStream().subscribe(output2::add);
    flowable._subscriber.onNext(new RecordEnvelope<>(new byte[] { 1, 1, 1 }));
    Assert.assertTrue(output0.poll() instanceof RecordEnvelope);
    Assert.assertTrue(output1.poll() instanceof RecordEnvelope);
    Assert.assertTrue(output2.poll() instanceof RecordEnvelope);
    flowable._subscriber.onNext(new RecordEnvelope<>(new byte[] { 1, 0, 0 }));
    Assert.assertTrue(output0.poll() instanceof RecordEnvelope);
    Assert.assertNull(output1.poll());
    Assert.assertNull(output2.poll());
    flowable._subscriber.onNext(new RecordEnvelope<>(new byte[] { 0, 1, 1 }));
    Assert.assertNull(output0.poll());
    Assert.assertTrue(output1.poll() instanceof RecordEnvelope);
    Assert.assertTrue(output2.poll() instanceof RecordEnvelope);
    flowable._subscriber.onNext(new BasicTestControlMessage<byte[]>("control"));
    Assert.assertTrue(output0.poll() instanceof BasicTestControlMessage);
    Assert.assertTrue(output1.poll() instanceof BasicTestControlMessage);
    Assert.assertTrue(output2.poll() instanceof BasicTestControlMessage);
    flowable._subscriber.onComplete();
}