Examples with StreamOperatorSpec org.apache.samza.operators.spec.StreamOperatorSpec used on opensource projects

Example 1 with StreamOperatorSpec

use of org.apache.samza.operators.spec.StreamOperatorSpec in project samza by apache.

the class TestMessageStreamImpl method testMerge.

@Test
public void testMerge() {
    StreamApplicationDescriptorImpl mockGraph = mock(StreamApplicationDescriptorImpl.class);
    OperatorSpec mockOpSpec1 = mock(OperatorSpec.class);
    MessageStream<TestMessageEnvelope> inputStream = new MessageStreamImpl<>(mockGraph, mockOpSpec1);
    // other streams have the same message type T as input stream message type M
    OperatorSpec mockOpSpec2 = mock(OperatorSpec.class);
    OperatorSpec mockOpSpec3 = mock(OperatorSpec.class);
    Collection<MessageStream<TestMessageEnvelope>> otherStreams1 = ImmutableList.of(new MessageStreamImpl<>(mockGraph, mockOpSpec2), new MessageStreamImpl<>(mockGraph, mockOpSpec3));
    inputStream.merge(otherStreams1);
    ArgumentCaptor<OperatorSpec> registeredOpCaptor1 = ArgumentCaptor.forClass(OperatorSpec.class);
    verify(mockOpSpec1).registerNextOperatorSpec(registeredOpCaptor1.capture());
    OperatorSpec<?, TestMessageEnvelope> registeredOpSpec1 = registeredOpCaptor1.getValue();
    assertTrue(registeredOpSpec1 instanceof StreamOperatorSpec);
    FlatMapFunction transformFn = ((StreamOperatorSpec) registeredOpSpec1).getTransformFn();
    ArgumentCaptor<OperatorSpec> registeredOpCaptor2 = ArgumentCaptor.forClass(OperatorSpec.class);
    verify(mockOpSpec2).registerNextOperatorSpec(registeredOpCaptor2.capture());
    OperatorSpec<?, TestMessageEnvelope> registeredOpSpec2 = registeredOpCaptor2.getValue();
    ArgumentCaptor<OperatorSpec> registeredOpCaptor3 = ArgumentCaptor.forClass(OperatorSpec.class);
    verify(mockOpSpec3).registerNextOperatorSpec(registeredOpCaptor3.capture());
    OperatorSpec<?, TestMessageEnvelope> registeredOpSpec3 = registeredOpCaptor3.getValue();
    assertEquals(registeredOpSpec1, registeredOpSpec2);
    assertEquals(registeredOpSpec2, registeredOpSpec3);
    assertEquals(OpCode.MERGE, registeredOpSpec1.getOpCode());
    assertNotNull(transformFn);
    TestMessageEnvelope mockInput = mock(TestMessageEnvelope.class);
    assertTrue(transformFn.apply(mockInput).contains(mockInput));
    assertEquals(1, transformFn.apply(mockInput).size());
}

Example 2 with StreamOperatorSpec

use of org.apache.samza.operators.spec.StreamOperatorSpec in project samza by apache.

the class TestMessageStreamImpl method testFlatMapWithRelaxedTypes.

@Test
public void testFlatMapWithRelaxedTypes() {
    StreamApplicationDescriptorImpl mockGraph = mock(StreamApplicationDescriptorImpl.class);
    OperatorSpec mockOpSpec = mock(OperatorSpec.class);
    MessageStreamImpl<TestInputMessageEnvelope> inputStream = new MessageStreamImpl<>(mockGraph, mockOpSpec);
    FlatMapFunction flatMapFunction = (FlatMapFunction<TestMessageEnvelope, TestOutputMessageEnvelope>) message -> Collections.emptyList();
    // should compile since TestInputMessageEnvelope extends TestMessageEnvelope
    inputStream.flatMap(flatMapFunction);
    ArgumentCaptor<OperatorSpec> registeredOpCaptor = ArgumentCaptor.forClass(OperatorSpec.class);
    verify(mockOpSpec).registerNextOperatorSpec(registeredOpCaptor.capture());
    OperatorSpec<?, TestMessageEnvelope> registeredOpSpec = registeredOpCaptor.getValue();
    assertTrue(registeredOpSpec instanceof StreamOperatorSpec);
    assertNotNull(((StreamOperatorSpec) registeredOpSpec).getTransformFn());
    assertEquals(OpCode.FLAT_MAP, registeredOpSpec.getOpCode());
}

Example 3 with StreamOperatorSpec

use of org.apache.samza.operators.spec.StreamOperatorSpec in project samza by apache.

the class TestFilterTranslator method testTranslate.

@Test
public void testTranslate() throws IOException, ClassNotFoundException {
    // setup mock values to the constructor of FilterTranslator
    LogicalFilter mockFilter = PowerMockito.mock(LogicalFilter.class);
    Context mockContext = mock(Context.class);
    ContainerContext mockContainerContext = mock(ContainerContext.class);
    TranslatorContext mockTranslatorContext = mock(TranslatorContext.class);
    TestMetricsRegistryImpl metricsRegistry = new TestMetricsRegistryImpl();
    RelNode mockInput = mock(RelNode.class);
    when(mockFilter.getInput()).thenReturn(mockInput);
    when(mockInput.getId()).thenReturn(1);
    when(mockFilter.getId()).thenReturn(2);
    StreamApplicationDescriptorImpl mockGraph = mock(StreamApplicationDescriptorImpl.class);
    OperatorSpec<Object, SamzaSqlRelMessage> mockInputOp = mock(OperatorSpec.class);
    MessageStream<SamzaSqlRelMessage> mockStream = new MessageStreamImpl<>(mockGraph, mockInputOp);
    when(mockTranslatorContext.getMessageStream(eq(1))).thenReturn(mockStream);
    doAnswer(this.getRegisterMessageStreamAnswer()).when(mockTranslatorContext).registerMessageStream(eq(2), any(MessageStream.class));
    RexToJavaCompiler mockCompiler = mock(RexToJavaCompiler.class);
    when(mockTranslatorContext.getExpressionCompiler()).thenReturn(mockCompiler);
    Expression mockExpr = mock(Expression.class);
    when(mockCompiler.compile(any(), any())).thenReturn(mockExpr);
    when(mockContext.getContainerContext()).thenReturn(mockContainerContext);
    when(mockContainerContext.getContainerMetricsRegistry()).thenReturn(metricsRegistry);
    // Apply translate() method to verify that we are getting the correct filter operator constructed
    FilterTranslator filterTranslator = new FilterTranslator(1);
    filterTranslator.translate(mockFilter, LOGICAL_OP_ID, mockTranslatorContext);
    // make sure that context has been registered with LogicFilter and output message streams
    verify(mockTranslatorContext, times(1)).registerRelNode(2, mockFilter);
    verify(mockTranslatorContext, times(1)).registerMessageStream(2, this.getRegisteredMessageStream(2));
    when(mockTranslatorContext.getRelNode(2)).thenReturn(mockFilter);
    when(mockTranslatorContext.getMessageStream(2)).thenReturn(this.getRegisteredMessageStream(2));
    StreamOperatorSpec filterSpec = (StreamOperatorSpec) Whitebox.getInternalState(this.getRegisteredMessageStream(2), "operatorSpec");
    assertNotNull(filterSpec);
    assertEquals(filterSpec.getOpCode(), OperatorSpec.OpCode.FILTER);
    // Verify that the describe() method will establish the context for the filter function
    Map<Integer, TranslatorContext> mockContexts = new HashMap<>();
    mockContexts.put(1, mockTranslatorContext);
    when(mockContext.getApplicationTaskContext()).thenReturn(new SamzaSqlApplicationContext(mockContexts));
    filterSpec.getTransformFn().init(mockContext);
    FilterFunction filterFn = (FilterFunction) Whitebox.getInternalState(filterSpec, "filterFn");
    assertNotNull(filterFn);
    assertEquals(mockTranslatorContext, Whitebox.getInternalState(filterFn, "translatorContext"));
    assertEquals(mockFilter, Whitebox.getInternalState(filterFn, "filter"));
    assertEquals(mockExpr, Whitebox.getInternalState(filterFn, "expr"));
    // Verify MetricsRegistry works with Project
    assertEquals(1, metricsRegistry.getGauges().size());
    assertTrue(metricsRegistry.getGauges().get(LOGICAL_OP_ID).size() > 0);
    assertEquals(1, metricsRegistry.getCounters().size());
    assertEquals(3, metricsRegistry.getCounters().get(LOGICAL_OP_ID).size());
    assertEquals(0, metricsRegistry.getCounters().get(LOGICAL_OP_ID).get(0).getCount());
    assertEquals(0, metricsRegistry.getCounters().get(LOGICAL_OP_ID).get(1).getCount());
    // Calling filterFn.apply() to verify the filter function is correctly applied to the input message
    SamzaSqlRelMessage mockInputMsg = new SamzaSqlRelMessage(new ArrayList<>(), new ArrayList<>(), new SamzaSqlRelMsgMetadata(0L, 0L));
    SamzaSqlExecutionContext executionContext = mock(SamzaSqlExecutionContext.class);
    DataContext dataContext = mock(DataContext.class);
    when(mockTranslatorContext.getExecutionContext()).thenReturn(executionContext);
    when(mockTranslatorContext.getDataContext()).thenReturn(dataContext);
    Object[] result = new Object[1];
    doAnswer(invocation -> {
        Object[] retValue = invocation.getArgumentAt(4, Object[].class);
        retValue[0] = new Boolean(true);
        return null;
    }).when(mockExpr).execute(eq(executionContext), eq(mockContext), eq(dataContext), eq(mockInputMsg.getSamzaSqlRelRecord().getFieldValues().toArray()), eq(result));
    assertTrue(filterFn.apply(mockInputMsg));
    doAnswer(invocation -> {
        Object[] retValue = invocation.getArgumentAt(4, Object[].class);
        retValue[0] = new Boolean(false);
        return null;
    }).when(mockExpr).execute(eq(executionContext), eq(mockContext), eq(dataContext), eq(mockInputMsg.getSamzaSqlRelRecord().getFieldValues().toArray()), eq(result));
    assertFalse(filterFn.apply(mockInputMsg));
    // Verify filterFn.apply() updates the MetricsRegistry metrics
    assertEquals(2, metricsRegistry.getCounters().get(LOGICAL_OP_ID).get(0).getCount());
    assertEquals(1, metricsRegistry.getCounters().get(LOGICAL_OP_ID).get(1).getCount());
}

Example 4 with StreamOperatorSpec

use of org.apache.samza.operators.spec.StreamOperatorSpec 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 5 with StreamOperatorSpec

use of org.apache.samza.operators.spec.StreamOperatorSpec in project samza by apache.

the class TestMessageStreamImpl method validateMergeOperator.

private void validateMergeOperator(MessageStream<TestMessageEnvelope> mergeSource, MessageStream<TestMessageEnvelope> mergeOutput) {
    Collection<OperatorSpec> subs = ((MessageStreamImpl<TestMessageEnvelope>) mergeSource).getRegisteredOperatorSpecs();
    assertEquals(subs.size(), 1);
    OperatorSpec<TestMessageEnvelope> mergeOp = subs.iterator().next();
    assertTrue(mergeOp instanceof StreamOperatorSpec);
    assertEquals(((StreamOperatorSpec) mergeOp).getNextStream(), mergeOutput);
    TestMessageEnvelope mockMsg = mock(TestMessageEnvelope.class);
    Collection<TestMessageEnvelope> outputs = ((StreamOperatorSpec<TestMessageEnvelope, TestMessageEnvelope>) mergeOp).getTransformFn().apply(mockMsg);
    assertEquals(outputs.size(), 1);
    assertEquals(outputs.iterator().next(), mockMsg);
}