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

Example 11 with TaskContext

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

the class TestWindowOperator method testTumblingWindowsDiscardingMode.

@Test
public void testTumblingWindowsDiscardingMode() throws Exception {
    StreamApplication sgb = new KeyedTumblingWindowStreamApplication(AccumulationMode.DISCARDING, Duration.ofSeconds(1), Triggers.repeat(Triggers.count(2)));
    List<WindowPane<Integer, Collection<IntegerEnvelope>>> windowPanes = new ArrayList<>();
    TestClock testClock = new TestClock();
    StreamOperatorTask task = new StreamOperatorTask(sgb, runner, testClock);
    task.init(config, taskContext);
    MessageCollector messageCollector = envelope -> windowPanes.add((WindowPane<Integer, Collection<IntegerEnvelope>>) envelope.getMessage());
    integers.forEach(n -> task.process(new IntegerEnvelope(n), messageCollector, taskCoordinator));
    testClock.advanceTime(Duration.ofSeconds(1));
    task.window(messageCollector, taskCoordinator);
    Assert.assertEquals(windowPanes.size(), 5);
    Assert.assertEquals(windowPanes.get(0).getKey().getKey(), new Integer(1));
    Assert.assertEquals((windowPanes.get(0).getMessage()).size(), 2);
    Assert.assertEquals(windowPanes.get(1).getKey().getKey(), new Integer(2));
    Assert.assertEquals((windowPanes.get(1).getMessage()).size(), 2);
    Assert.assertEquals(windowPanes.get(2).getKey().getKey(), new Integer(1));
    Assert.assertEquals((windowPanes.get(2).getMessage()).size(), 2);
    Assert.assertEquals(windowPanes.get(3).getKey().getKey(), new Integer(2));
    Assert.assertEquals((windowPanes.get(3).getMessage()).size(), 2);
    Assert.assertEquals(windowPanes.get(4).getKey().getKey(), new Integer(3));
    Assert.assertEquals((windowPanes.get(4).getMessage()).size(), 1);
}

Example 12 with TaskContext

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

the class TestWindowOperator method testSessionWindowsDiscardingMode.

@Test
public void testSessionWindowsDiscardingMode() throws Exception {
    StreamApplication sgb = new KeyedSessionWindowStreamApplication(AccumulationMode.DISCARDING, Duration.ofMillis(500));
    TestClock testClock = new TestClock();
    List<WindowPane<Integer, Collection<IntegerEnvelope>>> windowPanes = new ArrayList<>();
    StreamOperatorTask task = new StreamOperatorTask(sgb, runner, testClock);
    task.init(config, taskContext);
    MessageCollector messageCollector = envelope -> windowPanes.add((WindowPane<Integer, Collection<IntegerEnvelope>>) envelope.getMessage());
    task.process(new IntegerEnvelope(1), messageCollector, taskCoordinator);
    task.process(new IntegerEnvelope(1), messageCollector, taskCoordinator);
    testClock.advanceTime(Duration.ofSeconds(1));
    task.window(messageCollector, taskCoordinator);
    Assert.assertEquals(windowPanes.size(), 1);
    Assert.assertEquals(windowPanes.get(0).getKey().getPaneId(), "1");
    Assert.assertEquals(windowPanes.get(0).getKey().getKey(), new Integer(1));
    task.process(new IntegerEnvelope(2), messageCollector, taskCoordinator);
    task.process(new IntegerEnvelope(2), messageCollector, taskCoordinator);
    task.process(new IntegerEnvelope(3), messageCollector, taskCoordinator);
    task.process(new IntegerEnvelope(3), messageCollector, taskCoordinator);
    testClock.advanceTime(Duration.ofSeconds(1));
    task.window(messageCollector, taskCoordinator);
    Assert.assertEquals(windowPanes.size(), 3);
    Assert.assertEquals(windowPanes.get(0).getKey().getPaneId(), "1");
    Assert.assertEquals(windowPanes.get(1).getKey().getPaneId(), "1001");
    Assert.assertEquals(windowPanes.get(2).getKey().getPaneId(), "1001");
    Assert.assertEquals((windowPanes.get(0).getMessage()).size(), 2);
    Assert.assertEquals((windowPanes.get(1).getMessage()).size(), 2);
    Assert.assertEquals((windowPanes.get(2).getMessage()).size(), 2);
    task.process(new IntegerEnvelope(2), messageCollector, taskCoordinator);
    task.process(new IntegerEnvelope(2), messageCollector, taskCoordinator);
    testClock.advanceTime(Duration.ofSeconds(1));
    task.window(messageCollector, taskCoordinator);
    Assert.assertEquals(windowPanes.size(), 4);
    Assert.assertEquals(windowPanes.get(3).getKey().getKey(), new Integer(2));
    Assert.assertEquals(windowPanes.get(3).getKey().getPaneId(), "2001");
    Assert.assertEquals((windowPanes.get(3).getMessage()).size(), 2);
}

Example 13 with TaskContext

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

the class TestOperatorImpls method testEmptyChain.

@Test
public void testEmptyChain() throws InvocationTargetException, IllegalAccessException {
    // test creation of empty chain
    MessageStreamImpl<TestMessageEnvelope> testStream = mock(MessageStreamImpl.class);
    TaskContext mockContext = mock(TaskContext.class);
    when(mockContext.getMetricsRegistry()).thenReturn(new MetricsRegistryMap());
    Config mockConfig = mock(Config.class);
    OperatorImplGraph opGraph = new OperatorImplGraph();
    RootOperatorImpl operatorChain = (RootOperatorImpl) createOpsMethod.invoke(opGraph, testStream, mockConfig, mockContext);
    assertTrue(operatorChain != null);
}

Example 14 with TaskContext

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

the class TestOperatorImpls method testCreateOperator.

@Test
public void testCreateOperator() throws NoSuchFieldException, IllegalAccessException, InvocationTargetException {
    // get window operator
    WindowOperatorSpec mockWnd = mock(WindowOperatorSpec.class);
    WindowInternal<TestMessageEnvelope, String, Integer> windowInternal = new WindowInternal<>(null, null, null, null, null, WindowType.TUMBLING);
    when(mockWnd.getWindow()).thenReturn(windowInternal);
    Config mockConfig = mock(Config.class);
    TaskContext mockContext = mock(TaskContext.class);
    OperatorImplGraph opGraph = new OperatorImplGraph();
    OperatorImpl<TestMessageEnvelope, ?> opImpl = (OperatorImpl<TestMessageEnvelope, ?>) createOpMethod.invoke(opGraph, mockWnd, mockConfig, mockContext);
    assertTrue(opImpl instanceof WindowOperatorImpl);
    Field wndInternalField = WindowOperatorImpl.class.getDeclaredField("window");
    wndInternalField.setAccessible(true);
    WindowInternal wndInternal = (WindowInternal) wndInternalField.get(opImpl);
    assertEquals(wndInternal, windowInternal);
    // get simple operator
    StreamOperatorSpec<TestMessageEnvelope, TestOutputMessageEnvelope> mockSimpleOp = mock(StreamOperatorSpec.class);
    FlatMapFunction<TestMessageEnvelope, TestOutputMessageEnvelope> mockTxfmFn = mock(FlatMapFunction.class);
    when(mockSimpleOp.getTransformFn()).thenReturn(mockTxfmFn);
    opImpl = (OperatorImpl<TestMessageEnvelope, ?>) createOpMethod.invoke(opGraph, mockSimpleOp, mockConfig, mockContext);
    assertTrue(opImpl instanceof StreamOperatorImpl);
    Field txfmFnField = StreamOperatorImpl.class.getDeclaredField("transformFn");
    txfmFnField.setAccessible(true);
    assertEquals(mockTxfmFn, txfmFnField.get(opImpl));
    // get sink operator
    SinkFunction<TestMessageEnvelope> sinkFn = (m, mc, tc) -> {
    };
    SinkOperatorSpec<TestMessageEnvelope> sinkOp = mock(SinkOperatorSpec.class);
    when(sinkOp.getSinkFn()).thenReturn(sinkFn);
    opImpl = (OperatorImpl<TestMessageEnvelope, ?>) createOpMethod.invoke(opGraph, sinkOp, mockConfig, mockContext);
    assertTrue(opImpl instanceof SinkOperatorImpl);
    Field sinkFnField = SinkOperatorImpl.class.getDeclaredField("sinkFn");
    sinkFnField.setAccessible(true);
    assertEquals(sinkFn, sinkFnField.get(opImpl));
    // get join operator
    PartialJoinOperatorSpec<String, TestMessageEnvelope, TestMessageEnvelope, TestOutputMessageEnvelope> joinOp = mock(PartialJoinOperatorSpec.class);
    opImpl = (OperatorImpl<TestMessageEnvelope, ?>) createOpMethod.invoke(opGraph, joinOp, mockConfig, mockContext);
    assertTrue(opImpl instanceof PartialJoinOperatorImpl);
}

Example 15 with TaskContext

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

the class TestWindowOperator method testNonKeyedTumblingWindowsDiscardingMode.

@Test
public void testNonKeyedTumblingWindowsDiscardingMode() throws Exception {
    StreamApplication sgb = new TumblingWindowStreamApplication(AccumulationMode.DISCARDING, Duration.ofSeconds(1), Triggers.repeat(Triggers.count(1000)));
    List<WindowPane<Integer, Collection<IntegerEnvelope>>> windowPanes = new ArrayList<>();
    TestClock testClock = new TestClock();
    StreamOperatorTask task = new StreamOperatorTask(sgb, runner, testClock);
    task.init(config, taskContext);
    MessageCollector messageCollector = envelope -> windowPanes.add((WindowPane<Integer, Collection<IntegerEnvelope>>) envelope.getMessage());
    Assert.assertEquals(windowPanes.size(), 0);
    integers.forEach(n -> task.process(new IntegerEnvelope(n), messageCollector, taskCoordinator));
    Assert.assertEquals(windowPanes.size(), 0);
    testClock.advanceTime(Duration.ofSeconds(1));
    Assert.assertEquals(windowPanes.size(), 0);
    task.window(messageCollector, taskCoordinator);
    Assert.assertEquals(windowPanes.size(), 1);
    Assert.assertEquals((windowPanes.get(0).getMessage()).size(), 9);
}