Examples with TaskContext org.apache.samza.task.TaskContext used on opensource projects

Example 1 with TaskContext

use of org.apache.samza.task.TaskContext in project samza by apache.

the class TestJoinOperator method createStreamOperatorTask.

private StreamOperatorTask createStreamOperatorTask(Clock clock, StreamApplication app) throws Exception {
    ApplicationRunner runner = mock(ApplicationRunner.class);
    when(runner.getStreamSpec("instream")).thenReturn(new StreamSpec("instream", "instream", "insystem"));
    when(runner.getStreamSpec("instream2")).thenReturn(new StreamSpec("instream2", "instream2", "insystem2"));
    TaskContext taskContext = mock(TaskContext.class);
    when(taskContext.getSystemStreamPartitions()).thenReturn(ImmutableSet.of(new SystemStreamPartition("insystem", "instream", new Partition(0)), new SystemStreamPartition("insystem2", "instream2", new Partition(0))));
    when(taskContext.getMetricsRegistry()).thenReturn(new MetricsRegistryMap());
    Config config = mock(Config.class);
    StreamOperatorTask sot = new StreamOperatorTask(app, runner, clock);
    sot.init(config, taskContext);
    return sot;
}

Example 2 with TaskContext

use of org.apache.samza.task.TaskContext 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 3 with TaskContext

use of org.apache.samza.task.TaskContext in project samza by apache.

the class TestSinkOperatorImpl method createSinkOperator.

private SinkOperatorImpl createSinkOperator(SinkFunction<TestOutputMessageEnvelope> sinkFn) {
    SinkOperatorSpec<TestOutputMessageEnvelope> sinkOp = mock(SinkOperatorSpec.class);
    when(sinkOp.getSinkFn()).thenReturn(sinkFn);
    Config mockConfig = mock(Config.class);
    TaskContext mockContext = mock(TaskContext.class);
    return new SinkOperatorImpl<>(sinkOp, mockConfig, mockContext);
}

Example 4 with TaskContext

use of org.apache.samza.task.TaskContext 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);
}

Example 5 with TaskContext

use of org.apache.samza.task.TaskContext in project samza by apache.

the class TestWindowOperator method setup.

@Before
public void setup() throws Exception {
    config = mock(Config.class);
    taskContext = mock(TaskContext.class);
    runner = mock(ApplicationRunner.class);
    when(taskContext.getSystemStreamPartitions()).thenReturn(ImmutableSet.of(new SystemStreamPartition("kafka", "integers", new Partition(0))));
    when(taskContext.getMetricsRegistry()).thenReturn(new MetricsRegistryMap());
    when(runner.getStreamSpec("integers")).thenReturn(new StreamSpec("integers", "integers", "kafka"));
}