Examples with StreamOperatorSpec org.apache.samza.operators.spec.StreamOperatorSpec used on opensource projects

Example 6 with StreamOperatorSpec

use of org.apache.samza.operators.spec.StreamOperatorSpec in project samza by apache.

the class TestMessageStreamImpl method testFlatMap.

@Test
public void testFlatMap() {
    MessageStreamImpl<TestMessageEnvelope> inputStream = new MessageStreamImpl<>(mockGraph);
    List<TestOutputMessageEnvelope> flatOuts = new ArrayList<TestOutputMessageEnvelope>() {

        {
            this.add(mock(TestOutputMessageEnvelope.class));
            this.add(mock(TestOutputMessageEnvelope.class));
            this.add(mock(TestOutputMessageEnvelope.class));
        }
    };
    final List<TestMessageEnvelope> inputMsgs = new ArrayList<>();
    FlatMapFunction<TestMessageEnvelope, TestOutputMessageEnvelope> xFlatMap = (TestMessageEnvelope message) -> {
        inputMsgs.add(message);
        return flatOuts;
    };
    MessageStream<TestOutputMessageEnvelope> outputStream = inputStream.flatMap(xFlatMap);
    Collection<OperatorSpec> subs = inputStream.getRegisteredOperatorSpecs();
    assertEquals(subs.size(), 1);
    OperatorSpec<TestOutputMessageEnvelope> flatMapOp = subs.iterator().next();
    assertTrue(flatMapOp instanceof StreamOperatorSpec);
    assertEquals(flatMapOp.getNextStream(), outputStream);
    assertEquals(((StreamOperatorSpec) flatMapOp).getTransformFn(), xFlatMap);
    TestMessageEnvelope mockInput = mock(TestMessageEnvelope.class);
    // assert that the transformation function is what we defined above
    List<TestOutputMessageEnvelope> result = (List<TestOutputMessageEnvelope>) ((StreamOperatorSpec<TestMessageEnvelope, TestOutputMessageEnvelope>) flatMapOp).getTransformFn().apply(mockInput);
    assertEquals(flatOuts, result);
    assertEquals(inputMsgs.size(), 1);
    assertEquals(inputMsgs.get(0), mockInput);
}