Search in sources :

Example 26 with MessageCollector

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());
}
Also used : SystemStream(org.apache.samza.system.SystemStream) ArrayList(java.util.ArrayList) TaskCoordinator(org.apache.samza.task.TaskCoordinator) StreamOperatorSpec(org.apache.samza.operators.spec.StreamOperatorSpec) OperatorSpec(org.apache.samza.operators.spec.OperatorSpec) PartialJoinOperatorSpec(org.apache.samza.operators.spec.PartialJoinOperatorSpec) SinkOperatorSpec(org.apache.samza.operators.spec.SinkOperatorSpec) TestMessageEnvelope(org.apache.samza.operators.data.TestMessageEnvelope) MessageCollector(org.apache.samza.task.MessageCollector) SinkOperatorSpec(org.apache.samza.operators.spec.SinkOperatorSpec) OutgoingMessageEnvelope(org.apache.samza.system.OutgoingMessageEnvelope) MessageType(org.apache.samza.operators.data.MessageType) Test(org.junit.Test)

Example 27 with MessageCollector

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());
}
Also used : StreamOperatorTask(org.apache.samza.task.StreamOperatorTask) SystemStreamPartition(org.apache.samza.system.SystemStreamPartition) ArrayList(java.util.ArrayList) MessageCollector(org.apache.samza.task.MessageCollector) SystemStream(org.apache.samza.system.SystemStream) Duration(java.time.Duration) ApplicationRunner(org.apache.samza.runtime.ApplicationRunner) TaskContext(org.apache.samza.task.TaskContext) ImmutableSet(com.google.common.collect.ImmutableSet) IncomingMessageEnvelope(org.apache.samza.system.IncomingMessageEnvelope) Partition(org.apache.samza.Partition) Set(java.util.Set) Assert.assertTrue(org.junit.Assert.assertTrue) StreamSpec(org.apache.samza.system.StreamSpec) Clock(org.apache.samza.util.Clock) Test(org.junit.Test) Mockito.when(org.mockito.Mockito.when) JoinFunction(org.apache.samza.operators.functions.JoinFunction) TaskCoordinator(org.apache.samza.task.TaskCoordinator) List(java.util.List) SystemClock(org.apache.samza.util.SystemClock) TestClock(org.apache.samza.testUtils.TestClock) OutgoingMessageEnvelope(org.apache.samza.system.OutgoingMessageEnvelope) Config(org.apache.samza.config.Config) StreamApplication(org.apache.samza.application.StreamApplication) Assert.assertEquals(org.junit.Assert.assertEquals) MetricsRegistryMap(org.apache.samza.metrics.MetricsRegistryMap) Mockito.mock(org.mockito.Mockito.mock) SystemClock(org.apache.samza.util.SystemClock) StreamOperatorTask(org.apache.samza.task.StreamOperatorTask) MessageCollector(org.apache.samza.task.MessageCollector) ArrayList(java.util.ArrayList) Test(org.junit.Test)

Example 28 with MessageCollector

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);
}
Also used : StreamOperatorTask(org.apache.samza.task.StreamOperatorTask) SystemStreamPartition(org.apache.samza.system.SystemStreamPartition) ArrayList(java.util.ArrayList) MessageCollector(org.apache.samza.task.MessageCollector) SystemStream(org.apache.samza.system.SystemStream) Duration(java.time.Duration) ApplicationRunner(org.apache.samza.runtime.ApplicationRunner) TaskContext(org.apache.samza.task.TaskContext) ImmutableSet(com.google.common.collect.ImmutableSet) IncomingMessageEnvelope(org.apache.samza.system.IncomingMessageEnvelope) Partition(org.apache.samza.Partition) Set(java.util.Set) Assert.assertTrue(org.junit.Assert.assertTrue) StreamSpec(org.apache.samza.system.StreamSpec) Clock(org.apache.samza.util.Clock) Test(org.junit.Test) Mockito.when(org.mockito.Mockito.when) JoinFunction(org.apache.samza.operators.functions.JoinFunction) TaskCoordinator(org.apache.samza.task.TaskCoordinator) List(java.util.List) SystemClock(org.apache.samza.util.SystemClock) TestClock(org.apache.samza.testUtils.TestClock) OutgoingMessageEnvelope(org.apache.samza.system.OutgoingMessageEnvelope) Config(org.apache.samza.config.Config) StreamApplication(org.apache.samza.application.StreamApplication) Assert.assertEquals(org.junit.Assert.assertEquals) MetricsRegistryMap(org.apache.samza.metrics.MetricsRegistryMap) Mockito.mock(org.mockito.Mockito.mock) SystemClock(org.apache.samza.util.SystemClock) StreamOperatorTask(org.apache.samza.task.StreamOperatorTask) MessageCollector(org.apache.samza.task.MessageCollector) ArrayList(java.util.ArrayList) Test(org.junit.Test)

Aggregations

MessageCollector (org.apache.samza.task.MessageCollector)28 Test (org.junit.Test)28 TaskCoordinator (org.apache.samza.task.TaskCoordinator)27 Config (org.apache.samza.config.Config)23 TaskContext (org.apache.samza.task.TaskContext)23 ArrayList (java.util.ArrayList)21 MetricsRegistryMap (org.apache.samza.metrics.MetricsRegistryMap)20 OutgoingMessageEnvelope (org.apache.samza.system.OutgoingMessageEnvelope)20 SystemStream (org.apache.samza.system.SystemStream)20 StreamOperatorTask (org.apache.samza.task.StreamOperatorTask)19 ImmutableSet (com.google.common.collect.ImmutableSet)18 Duration (java.time.Duration)18 List (java.util.List)18 Partition (org.apache.samza.Partition)18 StreamApplication (org.apache.samza.application.StreamApplication)18 ApplicationRunner (org.apache.samza.runtime.ApplicationRunner)18 IncomingMessageEnvelope (org.apache.samza.system.IncomingMessageEnvelope)18 StreamSpec (org.apache.samza.system.StreamSpec)18 SystemStreamPartition (org.apache.samza.system.SystemStreamPartition)18 TestClock (org.apache.samza.testUtils.TestClock)18