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

Example 1 with TestMessageEnvelope

use of org.apache.samza.operators.data.TestMessageEnvelope 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);
}

Example 2 with TestMessageEnvelope

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

the class TestMessageStreamImpl method testJoin.

@Test
public void testJoin() {
    MessageStreamImpl<TestMessageEnvelope> source1 = new MessageStreamImpl<>(mockGraph);
    MessageStreamImpl<TestMessageEnvelope> source2 = new MessageStreamImpl<>(mockGraph);
    JoinFunction<String, TestMessageEnvelope, TestMessageEnvelope, TestOutputMessageEnvelope> joiner = new JoinFunction<String, TestMessageEnvelope, TestMessageEnvelope, TestOutputMessageEnvelope>() {

        @Override
        public TestOutputMessageEnvelope apply(TestMessageEnvelope m1, TestMessageEnvelope m2) {
            return new TestOutputMessageEnvelope(m1.getKey(), m1.getMessage().getValue().length() + m2.getMessage().getValue().length());
        }

        @Override
        public String getFirstKey(TestMessageEnvelope message) {
            return message.getKey();
        }

        @Override
        public String getSecondKey(TestMessageEnvelope message) {
            return message.getKey();
        }
    };
    MessageStream<TestOutputMessageEnvelope> joinOutput = source1.join(source2, joiner, Duration.ofMinutes(1));
    Collection<OperatorSpec> subs = source1.getRegisteredOperatorSpecs();
    assertEquals(subs.size(), 1);
    OperatorSpec<TestMessageEnvelope> joinOp1 = subs.iterator().next();
    assertTrue(joinOp1 instanceof PartialJoinOperatorSpec);
    assertEquals(((PartialJoinOperatorSpec) joinOp1).getNextStream(), joinOutput);
    subs = source2.getRegisteredOperatorSpecs();
    assertEquals(subs.size(), 1);
    OperatorSpec<TestMessageEnvelope> joinOp2 = subs.iterator().next();
    assertTrue(joinOp2 instanceof PartialJoinOperatorSpec);
    assertEquals(((PartialJoinOperatorSpec) joinOp2).getNextStream(), joinOutput);
    TestMessageEnvelope joinMsg1 = new TestMessageEnvelope("test-join-1", "join-msg-001", 11111L);
    TestMessageEnvelope joinMsg2 = new TestMessageEnvelope("test-join-2", "join-msg-002", 22222L);
    TestOutputMessageEnvelope xOut = (TestOutputMessageEnvelope) ((PartialJoinOperatorSpec) joinOp1).getThisPartialJoinFn().apply(joinMsg1, joinMsg2);
    assertEquals(xOut.getKey(), "test-join-1");
    assertEquals(xOut.getMessage(), Integer.valueOf(24));
    xOut = (TestOutputMessageEnvelope) ((PartialJoinOperatorSpec) joinOp2).getThisPartialJoinFn().apply(joinMsg2, joinMsg1);
    assertEquals(xOut.getKey(), "test-join-1");
    assertEquals(xOut.getMessage(), Integer.valueOf(24));
}

Example 3 with TestMessageEnvelope

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

the class TestMessageStreamImpl method testPartitionBy.

@Test
public void testPartitionBy() {
    Map<String, String> map = new HashMap<>();
    map.put(JobConfig.JOB_DEFAULT_SYSTEM(), "testsystem");
    Config config = new MapConfig(map);
    ApplicationRunner runner = ApplicationRunner.fromConfig(config);
    StreamGraphImpl streamGraph = new StreamGraphImpl(runner, config);
    MessageStreamImpl<TestMessageEnvelope> inputStream = new MessageStreamImpl<>(streamGraph);
    Function<TestMessageEnvelope, String> keyExtractorFunc = m -> "222";
    inputStream.partitionBy(keyExtractorFunc);
    assertTrue(streamGraph.getInputStreams().size() == 1);
    assertTrue(streamGraph.getOutputStreams().size() == 1);
    Collection<OperatorSpec> subs = inputStream.getRegisteredOperatorSpecs();
    assertEquals(subs.size(), 1);
    OperatorSpec<TestMessageEnvelope> partitionByOp = subs.iterator().next();
    assertTrue(partitionByOp instanceof SinkOperatorSpec);
    assertNull(partitionByOp.getNextStream());
    ((SinkOperatorSpec) partitionByOp).getSinkFn().apply(new TestMessageEnvelope("111", "test", 1000), envelope -> assertTrue(envelope.getPartitionKey().equals("222")), null);
}

Example 4 with TestMessageEnvelope

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

the class TestOperatorImpls method testJoinChain.

@Test
public void testJoinChain() throws IllegalAccessException, InvocationTargetException {
    // test creation of join chain
    StreamGraphImpl mockGraph = mock(StreamGraphImpl.class);
    MessageStreamImpl<TestMessageEnvelope> input1 = TestMessageStreamImplUtil.getMessageStreamImpl(mockGraph);
    MessageStreamImpl<TestMessageEnvelope> input2 = TestMessageStreamImplUtil.getMessageStreamImpl(mockGraph);
    TaskContext mockContext = mock(TaskContext.class);
    when(mockContext.getMetricsRegistry()).thenReturn(new MetricsRegistryMap());
    Config mockConfig = mock(Config.class);
    input1.join(input2, new JoinFunction<String, TestMessageEnvelope, TestMessageEnvelope, TestOutputMessageEnvelope>() {

        @Override
        public TestOutputMessageEnvelope apply(TestMessageEnvelope m1, TestMessageEnvelope m2) {
            return new TestOutputMessageEnvelope(m1.getKey(), m1.getMessage().getValue().length() + m2.getMessage().getValue().length());
        }

        @Override
        public String getFirstKey(TestMessageEnvelope message) {
            return message.getKey();
        }

        @Override
        public String getSecondKey(TestMessageEnvelope message) {
            return message.getKey();
        }
    }, Duration.ofMinutes(1)).map(m -> m);
    OperatorImplGraph opGraph = new OperatorImplGraph();
    // now, we create chained operators from each input sources
    RootOperatorImpl chain1 = (RootOperatorImpl) createOpsMethod.invoke(opGraph, input1, mockConfig, mockContext);
    RootOperatorImpl chain2 = (RootOperatorImpl) createOpsMethod.invoke(opGraph, input2, mockConfig, mockContext);
    // check that those two chains will merge at map operator
    // first branch of the join
    Set<OperatorImpl> subsSet = (Set<OperatorImpl>) nextOperatorsField.get(chain1);
    assertEquals(subsSet.size(), 1);
    OperatorImpl<TestMessageEnvelope, TestOutputMessageEnvelope> joinOp1 = subsSet.iterator().next();
    Set<OperatorImpl> subsOps = (Set<OperatorImpl>) nextOperatorsField.get(joinOp1);
    assertEquals(subsOps.size(), 1);
    // the map operator consumes the common join output, where two branches merge
    OperatorImpl mapImpl = subsOps.iterator().next();
    // second branch of the join
    subsSet = (Set<OperatorImpl>) nextOperatorsField.get(chain2);
    assertEquals(subsSet.size(), 1);
    OperatorImpl<TestMessageEnvelope, TestOutputMessageEnvelope> joinOp2 = subsSet.iterator().next();
    assertNotSame(joinOp1, joinOp2);
    subsOps = (Set<OperatorImpl>) nextOperatorsField.get(joinOp2);
    assertEquals(subsOps.size(), 1);
    // make sure that the map operator is the same
    assertEquals(mapImpl, subsOps.iterator().next());
}

Example 5 with TestMessageEnvelope

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

the class TestStreamOperatorImpl method testSimpleOperator.

@Test
@SuppressWarnings("unchecked")
public void testSimpleOperator() {
    StreamOperatorSpec<TestMessageEnvelope, TestOutputMessageEnvelope> mockOp = mock(StreamOperatorSpec.class);
    FlatMapFunction<TestMessageEnvelope, TestOutputMessageEnvelope> txfmFn = mock(FlatMapFunction.class);
    when(mockOp.getTransformFn()).thenReturn(txfmFn);
    Config mockConfig = mock(Config.class);
    TaskContext mockContext = mock(TaskContext.class);
    StreamOperatorImpl<TestMessageEnvelope, TestOutputMessageEnvelope> opImpl = spy(new StreamOperatorImpl<>(mockOp, mockConfig, mockContext));
    TestMessageEnvelope inMsg = mock(TestMessageEnvelope.class);
    TestOutputMessageEnvelope outMsg = mock(TestOutputMessageEnvelope.class);
    Collection<TestOutputMessageEnvelope> mockOutputs = new ArrayList() {

        {
            this.add(outMsg);
        }
    };
    when(txfmFn.apply(inMsg)).thenReturn(mockOutputs);
    MessageCollector mockCollector = mock(MessageCollector.class);
    TaskCoordinator mockCoordinator = mock(TaskCoordinator.class);
    Collection<TestOutputMessageEnvelope> results = opImpl.handleMessage(inMsg, mockCollector, mockCoordinator);
    verify(txfmFn, times(1)).apply(inMsg);
    assertEquals(results, mockOutputs);
}