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

Example 6 with OperatorSpec

use of org.apache.samza.operators.spec.OperatorSpec 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 7 with OperatorSpec

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

the class TestMessageStreamImpl method validateMergeOperator.

private void validateMergeOperator(MessageStream<TestMessageEnvelope> mergeSource, MessageStream<TestMessageEnvelope> mergeOutput) {
    Collection<OperatorSpec> subs = ((MessageStreamImpl<TestMessageEnvelope>) mergeSource).getRegisteredOperatorSpecs();
    assertEquals(subs.size(), 1);
    OperatorSpec<TestMessageEnvelope> mergeOp = subs.iterator().next();
    assertTrue(mergeOp instanceof StreamOperatorSpec);
    assertEquals(((StreamOperatorSpec) mergeOp).getNextStream(), mergeOutput);
    TestMessageEnvelope mockMsg = mock(TestMessageEnvelope.class);
    Collection<TestMessageEnvelope> outputs = ((StreamOperatorSpec<TestMessageEnvelope, TestMessageEnvelope>) mergeOp).getTransformFn().apply(mockMsg);
    assertEquals(outputs.size(), 1);
    assertEquals(outputs.iterator().next(), mockMsg);
}

Example 8 with OperatorSpec

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

the class TestMessageStreamImpl method testFilter.

@Test
public void testFilter() {
    MessageStreamImpl<TestMessageEnvelope> inputStream = new MessageStreamImpl<>(mockGraph);
    FilterFunction<TestMessageEnvelope> xFilter = (TestMessageEnvelope m) -> m.getMessage().getEventTime() > 123456L;
    MessageStream<TestMessageEnvelope> outputStream = inputStream.filter(xFilter);
    Collection<OperatorSpec> subs = inputStream.getRegisteredOperatorSpecs();
    assertEquals(subs.size(), 1);
    OperatorSpec<TestMessageEnvelope> filterOp = subs.iterator().next();
    assertTrue(filterOp instanceof StreamOperatorSpec);
    assertEquals(filterOp.getNextStream(), outputStream);
    // assert that the transformation function is what we defined above
    FlatMapFunction<TestMessageEnvelope, TestMessageEnvelope> txfmFn = ((StreamOperatorSpec<TestMessageEnvelope, TestMessageEnvelope>) filterOp).getTransformFn();
    TestMessageEnvelope mockMsg = mock(TestMessageEnvelope.class);
    MessageType mockInnerTestMessage = mock(MessageType.class);
    when(mockMsg.getMessage()).thenReturn(mockInnerTestMessage);
    when(mockInnerTestMessage.getEventTime()).thenReturn(11111L);
    Collection<TestMessageEnvelope> output = txfmFn.apply(mockMsg);
    assertTrue(output.isEmpty());
    when(mockMsg.getMessage()).thenReturn(mockInnerTestMessage);
    when(mockInnerTestMessage.getEventTime()).thenReturn(999999L);
    output = txfmFn.apply(mockMsg);
    assertEquals(output.size(), 1);
    assertEquals(output.iterator().next(), mockMsg);
}

Example 9 with OperatorSpec

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

the class TestMessageStreamImpl method testSink.

@Test
public void testSink() {
    MessageStreamImpl<TestMessageEnvelope> inputStream = new MessageStreamImpl<>(mockGraph);
    SystemStream testStream = new SystemStream("test-sys", "test-stream");
    SinkFunction<TestMessageEnvelope> xSink = (TestMessageEnvelope m, MessageCollector mc, TaskCoordinator tc) -> {
        mc.send(new OutgoingMessageEnvelope(testStream, m.getMessage()));
        tc.commit(TaskCoordinator.RequestScope.CURRENT_TASK);
    };
    inputStream.sink(xSink);
    Collection<OperatorSpec> subs = inputStream.getRegisteredOperatorSpecs();
    assertEquals(subs.size(), 1);
    OperatorSpec<TestMessageEnvelope> sinkOp = subs.iterator().next();
    assertTrue(sinkOp instanceof SinkOperatorSpec);
    assertEquals(((SinkOperatorSpec) sinkOp).getSinkFn(), xSink);
    TestMessageEnvelope mockTest1 = mock(TestMessageEnvelope.class);
    MessageType mockMsgBody = mock(MessageType.class);
    when(mockTest1.getMessage()).thenReturn(mockMsgBody);
    final List<OutgoingMessageEnvelope> outMsgs = new ArrayList<>();
    MessageCollector mockCollector = mock(MessageCollector.class);
    doAnswer(invocation -> {
        outMsgs.add((OutgoingMessageEnvelope) invocation.getArguments()[0]);
        return null;
    }).when(mockCollector).send(any());
    TaskCoordinator mockCoordinator = mock(TaskCoordinator.class);
    ((SinkOperatorSpec) sinkOp).getSinkFn().apply(mockTest1, mockCollector, mockCoordinator);
    assertEquals(1, outMsgs.size());
    assertEquals(testStream, outMsgs.get(0).getSystemStream());
    assertEquals(mockMsgBody, outMsgs.get(0).getMessage());
}

Example 10 with OperatorSpec

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

the class TestMessageStreamImpl method testMap.

@Test
public void testMap() {
    MessageStreamImpl<TestMessageEnvelope> inputStream = new MessageStreamImpl<>(mockGraph);
    MapFunction<TestMessageEnvelope, TestOutputMessageEnvelope> xMap = (TestMessageEnvelope m) -> new TestOutputMessageEnvelope(m.getKey(), m.getMessage().getValue().length() + 1);
    MessageStream<TestOutputMessageEnvelope> outputStream = inputStream.map(xMap);
    Collection<OperatorSpec> subs = inputStream.getRegisteredOperatorSpecs();
    assertEquals(subs.size(), 1);
    OperatorSpec<TestOutputMessageEnvelope> mapOp = subs.iterator().next();
    assertTrue(mapOp instanceof StreamOperatorSpec);
    assertEquals(mapOp.getNextStream(), outputStream);
    // assert that the transformation function is what we defined above
    TestMessageEnvelope xTestMsg = mock(TestMessageEnvelope.class);
    MessageType mockInnerTestMessage = mock(MessageType.class);
    when(xTestMsg.getKey()).thenReturn("test-msg-key");
    when(xTestMsg.getMessage()).thenReturn(mockInnerTestMessage);
    when(mockInnerTestMessage.getValue()).thenReturn("123456789");
    Collection<TestOutputMessageEnvelope> cOutputMsg = ((StreamOperatorSpec<TestMessageEnvelope, TestOutputMessageEnvelope>) mapOp).getTransformFn().apply(xTestMsg);
    assertEquals(cOutputMsg.size(), 1);
    TestOutputMessageEnvelope outputMessage = cOutputMsg.iterator().next();
    assertEquals(outputMessage.getKey(), xTestMsg.getKey());
    assertEquals(outputMessage.getMessage(), Integer.valueOf(xTestMsg.getMessage().getValue().length() + 1));
}