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

Example 21 with TaskContext

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

the class TestWindowOperator method testCancellationOfOnceTrigger.

@Test
public void testCancellationOfOnceTrigger() throws Exception {
    StreamApplication sgb = new KeyedTumblingWindowStreamApplication(AccumulationMode.ACCUMULATING, Duration.ofSeconds(1), Triggers.count(2));
    TestClock testClock = new TestClock();
    StreamOperatorTask task = new StreamOperatorTask(sgb, runner, testClock);
    task.init(config, taskContext);
    List<WindowPane<Integer, Collection<IntegerEnvelope>>> windowPanes = new ArrayList<>();
    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);
    Assert.assertEquals(windowPanes.size(), 1);
    Assert.assertEquals(windowPanes.get(0).getKey().getPaneId(), "0");
    Assert.assertEquals(windowPanes.get(0).getKey().getKey(), new Integer(1));
    Assert.assertEquals(windowPanes.get(0).getFiringType(), FiringType.EARLY);
    task.process(new IntegerEnvelope(1), messageCollector, taskCoordinator);
    task.process(new IntegerEnvelope(1), messageCollector, taskCoordinator);
    task.process(new IntegerEnvelope(1), messageCollector, taskCoordinator);
    Assert.assertEquals(windowPanes.size(), 1);
    testClock.advanceTime(Duration.ofSeconds(1));
    task.window(messageCollector, taskCoordinator);
    Assert.assertEquals(windowPanes.size(), 2);
    Assert.assertEquals(windowPanes.get(0).getKey().getPaneId(), "0");
    Assert.assertEquals(windowPanes.get(1).getKey().getPaneId(), "0");
    Assert.assertEquals(windowPanes.get(1).getFiringType(), FiringType.DEFAULT);
    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(2).getKey().getKey(), new Integer(3));
    Assert.assertEquals(windowPanes.get(2).getKey().getPaneId(), "1000");
    Assert.assertEquals(windowPanes.get(2).getFiringType(), FiringType.DEFAULT);
    Assert.assertEquals((windowPanes.get(2).getMessage()).size(), 1);
}

Example 22 with TaskContext

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

the class TestOperatorImpl method testMultipleInitShouldThrow.

@Test(expected = IllegalStateException.class)
public void testMultipleInitShouldThrow() {
    OperatorImpl<Object, Object> opImpl = new TestOpImpl(mock(Object.class));
    TaskContext mockTaskContext = mock(TaskContext.class);
    when(mockTaskContext.getMetricsRegistry()).thenReturn(new MetricsRegistryMap());
    opImpl.init(mock(Config.class), mockTaskContext);
    opImpl.init(mock(Config.class), mockTaskContext);
}

Example 23 with TaskContext

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

the class TestWindowOperator method testCancellationOfAnyTrigger.

@Test
public void testCancellationOfAnyTrigger() throws Exception {
    StreamApplication sgb = new KeyedTumblingWindowStreamApplication(AccumulationMode.ACCUMULATING, Duration.ofSeconds(1), Triggers.any(Triggers.count(2), Triggers.timeSinceFirstMessage(Duration.ofMillis(500))));
    TestClock testClock = new TestClock();
    StreamOperatorTask task = new StreamOperatorTask(sgb, runner, testClock);
    task.init(config, taskContext);
    List<WindowPane<Integer, Collection<IntegerEnvelope>>> windowPanes = new ArrayList<>();
    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);
    //assert that the count trigger fired
    Assert.assertEquals(windowPanes.size(), 1);
    //advance the timer to enable the triggering of the inner timeSinceFirstMessage trigger
    testClock.advanceTime(Duration.ofMillis(500));
    //assert that the triggering of the count trigger cancelled the inner timeSinceFirstMessage trigger
    Assert.assertEquals(windowPanes.size(), 1);
    task.process(new IntegerEnvelope(1), messageCollector, taskCoordinator);
    task.process(new IntegerEnvelope(1), messageCollector, taskCoordinator);
    task.process(new IntegerEnvelope(1), messageCollector, taskCoordinator);
    //advance timer by 500 more millis to enable the default trigger
    testClock.advanceTime(Duration.ofMillis(500));
    task.window(messageCollector, taskCoordinator);
    //assert that the default trigger fired
    Assert.assertEquals(windowPanes.size(), 2);
    Assert.assertEquals(windowPanes.get(1).getFiringType(), FiringType.DEFAULT);
    Assert.assertEquals(windowPanes.get(1).getKey().getKey(), new Integer(1));
    Assert.assertEquals(windowPanes.get(1).getKey().getPaneId(), "0");
    Assert.assertEquals((windowPanes.get(1).getMessage()).size(), 5);
    task.process(new IntegerEnvelope(1), messageCollector, taskCoordinator);
    //advance timer by 500 millis to enable the inner timeSinceFirstMessage trigger
    testClock.advanceTime(Duration.ofMillis(500));
    task.window(messageCollector, taskCoordinator);
    Assert.assertEquals(windowPanes.size(), 3);
    Assert.assertEquals(windowPanes.get(2).getFiringType(), FiringType.EARLY);
    Assert.assertEquals(windowPanes.get(2).getKey().getKey(), new Integer(1));
    Assert.assertEquals(windowPanes.get(2).getKey().getPaneId(), "1000");
    //advance timer by > 500 millis to enable the default trigger
    testClock.advanceTime(Duration.ofMillis(900));
    task.window(messageCollector, taskCoordinator);
    Assert.assertEquals(windowPanes.size(), 4);
    Assert.assertEquals(windowPanes.get(3).getFiringType(), FiringType.DEFAULT);
    Assert.assertEquals(windowPanes.get(3).getKey().getKey(), new Integer(1));
    Assert.assertEquals(windowPanes.get(3).getKey().getPaneId(), "1000");
}

Example 24 with TaskContext

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

the class TestStreamOperatorImpl method testSimpleOperatorClose.

@Test
public void testSimpleOperatorClose() {
    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));
    // ensure that close is not called yet
    verify(txfmFn, times(0)).close();
    opImpl.handleClose();
    // ensure that close is called once inside handleClose()
    verify(txfmFn, times(1)).close();
}