FilterTranslator { void translate(final LogicalFilter filter, final String logicalOpId, final TranslatorContext context) { MessageStream<SamzaSqlRelMessage> inputStream = context.getMessageStream(filter.getInput().getId()); final int filterId = filter.getId(); MessageStream<SamzaSqlRelMessage> outputStream = inputStream.filter(new FilterTranslatorFunction(filterId, queryId, logicalOpId)); context.registerMessageStream(filterId, outputStream); context.registerRelNode(filterId, filter); } FilterTranslator(int queryId);   }
@Test public void testTranslate() throws IOException, ClassNotFoundException { 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); FilterTranslator filterTranslator = new FilterTranslator(1); filterTranslator.translate(mockFilter, LOGICAL_OP_ID, mockTranslatorContext); 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); 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")); 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()); 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)); assertEquals(2, metricsRegistry.getCounters().get(LOGICAL_OP_ID).get(0).getCount()); assertEquals(1, metricsRegistry.getCounters().get(LOGICAL_OP_ID).get(1).getCount()); }