Examples with TestExtOutputMessageEnvelope org.apache.samza.operators.data.TestExtOutputMessageEnvelope used on opensource projects

Example 1 with TestExtOutputMessageEnvelope

use of org.apache.samza.operators.data.TestExtOutputMessageEnvelope in project samza by apache.

the class TestMessageStreamImpl method testFlatMapWithRelaxedTypes.

@Test
public void testFlatMapWithRelaxedTypes() {
    MessageStreamImpl<TestInputMessageEnvelope> inputStream = new MessageStreamImpl<>(mockGraph);
    List<TestExtOutputMessageEnvelope> flatOuts = new ArrayList<TestExtOutputMessageEnvelope>() {

        {
            this.add(new TestExtOutputMessageEnvelope("output-key-1", 1, "output-id-001"));
            this.add(new TestExtOutputMessageEnvelope("output-key-2", 2, "output-id-002"));
            this.add(new TestExtOutputMessageEnvelope("output-key-3", 3, "output-id-003"));
        }
    };
    class MyFlatMapFunction implements FlatMapFunction<TestMessageEnvelope, TestExtOutputMessageEnvelope> {

        public final List<TestMessageEnvelope> inputMsgs = new ArrayList<>();

        @Override
        public Collection<TestExtOutputMessageEnvelope> apply(TestMessageEnvelope message) {
            inputMsgs.add(message);
            return flatOuts;
        }

        @Override
        public void init(Config config, TaskContext context) {
            inputMsgs.clear();
        }
    }
    MyFlatMapFunction xFlatMap = new MyFlatMapFunction();
    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(xFlatMap.inputMsgs.size(), 1);
    assertEquals(xFlatMap.inputMsgs.get(0), mockInput);
}