use of org.apache.samza.task.MessageCollector 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());
}
use of org.apache.samza.task.MessageCollector in project samza by apache.
the class TestJoinOperator method joinNoMatch.
@Test
public void joinNoMatch() throws Exception {
StreamOperatorTask sot = createStreamOperatorTask(new SystemClock(), new TestJoinStreamApplication(new TestJoinFunction()));
List<Integer> output = new ArrayList<>();
MessageCollector messageCollector = envelope -> output.add((Integer) envelope.getMessage());
// push messages to first stream
numbers.forEach(n -> sot.process(new FirstStreamIME(n, n), messageCollector, taskCoordinator));
// push messages to second stream with different keys
numbers.forEach(n -> sot.process(new SecondStreamIME(n + 100, n), messageCollector, taskCoordinator));
assertTrue(output.isEmpty());
}
use of org.apache.samza.task.MessageCollector in project samza by apache.
the class TestJoinOperator method join.
@Test
public void join() throws Exception {
StreamOperatorTask sot = createStreamOperatorTask(new SystemClock(), new TestJoinStreamApplication(new TestJoinFunction()));
List<Integer> output = new ArrayList<>();
MessageCollector messageCollector = envelope -> output.add((Integer) envelope.getMessage());
// push messages to first stream
numbers.forEach(n -> sot.process(new FirstStreamIME(n, n), messageCollector, taskCoordinator));
// push messages to second stream with same keys
numbers.forEach(n -> sot.process(new SecondStreamIME(n, n), messageCollector, taskCoordinator));
int outputSum = output.stream().reduce(0, (s, m) -> s + m);
assertEquals(110, outputSum);
}
Aggregations