OperatorImplGraph { public InputOperatorImpl getInputOperator(SystemStream systemStream) { return this.inputOperators.get(systemStream); } OperatorImplGraph(OperatorSpecGraph specGraph, Context context, Clock clock); InputOperatorImpl getInputOperator(SystemStream systemStream); void close(); Collection<InputOperatorImpl> getAllInputOperators();  }
@Test public void testLinearChain() { String inputStreamId = "input"; String inputSystem = "input-system"; String inputPhysicalName = "input-stream"; String outputStreamId = "output"; String outputSystem = "output-system"; String outputPhysicalName = "output-stream"; String intermediateSystem = "intermediate-system"; HashMap<String, String> configs = new HashMap<>(); configs.put(JobConfig.JOB_NAME, "jobName"); configs.put(JobConfig.JOB_ID, "jobId"); configs.put(JobConfig.JOB_DEFAULT_SYSTEM, intermediateSystem); StreamTestUtils.addStreamConfigs(configs, inputStreamId, inputSystem, inputPhysicalName); StreamTestUtils.addStreamConfigs(configs, outputStreamId, outputSystem, outputPhysicalName); Config config = new MapConfig(configs); when(this.context.getJobContext().getConfig()).thenReturn(config); StreamApplicationDescriptorImpl graphSpec = new StreamApplicationDescriptorImpl(appDesc -> { GenericSystemDescriptor sd = new GenericSystemDescriptor(inputSystem, "mockFactoryClass"); GenericInputDescriptor inputDescriptor = sd.getInputDescriptor(inputStreamId, mock(Serde.class)); GenericOutputDescriptor outputDescriptor = sd.getOutputDescriptor(outputStreamId, mock(Serde.class)); MessageStream<Object> inputStream = appDesc.getInputStream(inputDescriptor); OutputStream<Object> outputStream = appDesc.getOutputStream(outputDescriptor); inputStream .filter(mock(FilterFunction.class)) .map(mock(MapFunction.class)) .sendTo(outputStream); }, config); OperatorImplGraph opImplGraph = new OperatorImplGraph(graphSpec.getOperatorSpecGraph(), this.context, mock(Clock.class)); InputOperatorImpl inputOpImpl = opImplGraph.getInputOperator(new SystemStream(inputSystem, inputPhysicalName)); assertEquals(1, inputOpImpl.registeredOperators.size()); OperatorImpl filterOpImpl = (FlatmapOperatorImpl) inputOpImpl.registeredOperators.iterator().next(); assertEquals(1, filterOpImpl.registeredOperators.size()); assertEquals(OpCode.FILTER, filterOpImpl.getOperatorSpec().getOpCode()); OperatorImpl mapOpImpl = (FlatmapOperatorImpl) filterOpImpl.registeredOperators.iterator().next(); assertEquals(1, mapOpImpl.registeredOperators.size()); assertEquals(OpCode.MAP, mapOpImpl.getOperatorSpec().getOpCode()); OperatorImpl sendToOpImpl = (OutputOperatorImpl) mapOpImpl.registeredOperators.iterator().next(); assertEquals(0, sendToOpImpl.registeredOperators.size()); assertEquals(OpCode.SEND_TO, sendToOpImpl.getOperatorSpec().getOpCode()); }
@Test public void testPartitionByChain() { String inputStreamId = "input"; String inputSystem = "input-system"; String inputPhysicalName = "input-stream"; String outputStreamId = "output"; String outputSystem = "output-system"; String outputPhysicalName = "output-stream"; String intermediateStreamId = "jobName-jobId-partition_by-p1"; String intermediateSystem = "intermediate-system"; HashMap<String, String> configs = new HashMap<>(); configs.put(JobConfig.JOB_NAME, "jobName"); configs.put(JobConfig.JOB_ID, "jobId"); configs.put(JobConfig.JOB_DEFAULT_SYSTEM, intermediateSystem); StreamTestUtils.addStreamConfigs(configs, inputStreamId, inputSystem, inputPhysicalName); StreamTestUtils.addStreamConfigs(configs, outputStreamId, outputSystem, outputPhysicalName); Config config = new MapConfig(configs); when(this.context.getJobContext().getConfig()).thenReturn(config); StreamApplicationDescriptorImpl graphSpec = new StreamApplicationDescriptorImpl(appDesc -> { GenericSystemDescriptor isd = new GenericSystemDescriptor(inputSystem, "mockFactoryClass"); GenericSystemDescriptor osd = new GenericSystemDescriptor(outputSystem, "mockFactoryClass"); GenericInputDescriptor inputDescriptor = isd.getInputDescriptor(inputStreamId, mock(Serde.class)); GenericOutputDescriptor outputDescriptor = osd.getOutputDescriptor(outputStreamId, KVSerde.of(mock(IntegerSerde.class), mock(StringSerde.class))); MessageStream<Object> inputStream = appDesc.getInputStream(inputDescriptor); OutputStream<KV<Integer, String>> outputStream = appDesc.getOutputStream(outputDescriptor); inputStream .partitionBy(Object::hashCode, Object::toString, KVSerde.of(mock(IntegerSerde.class), mock(StringSerde.class)), "p1") .sendTo(outputStream); }, config); JobModel jobModel = mock(JobModel.class); ContainerModel containerModel = mock(ContainerModel.class); TaskModel taskModel = mock(TaskModel.class); when(jobModel.getContainers()).thenReturn(Collections.singletonMap("0", containerModel)); when(containerModel.getTasks()).thenReturn(Collections.singletonMap(new TaskName("task 0"), taskModel)); when(taskModel.getSystemStreamPartitions()).thenReturn(Collections.emptySet()); when(((TaskContextImpl) this.context.getTaskContext()).getJobModel()).thenReturn(jobModel); OperatorImplGraph opImplGraph = new OperatorImplGraph(graphSpec.getOperatorSpecGraph(), this.context, mock(Clock.class)); InputOperatorImpl inputOpImpl = opImplGraph.getInputOperator(new SystemStream(inputSystem, inputPhysicalName)); assertEquals(1, inputOpImpl.registeredOperators.size()); OperatorImpl partitionByOpImpl = (PartitionByOperatorImpl) inputOpImpl.registeredOperators.iterator().next(); assertEquals(0, partitionByOpImpl.registeredOperators.size()); assertEquals(OpCode.PARTITION_BY, partitionByOpImpl.getOperatorSpec().getOpCode()); InputOperatorImpl repartitionedInputOpImpl = opImplGraph.getInputOperator(new SystemStream(intermediateSystem, intermediateStreamId)); assertEquals(1, repartitionedInputOpImpl.registeredOperators.size()); OperatorImpl sendToOpImpl = (OutputOperatorImpl) repartitionedInputOpImpl.registeredOperators.iterator().next(); assertEquals(0, sendToOpImpl.registeredOperators.size()); assertEquals(OpCode.SEND_TO, sendToOpImpl.getOperatorSpec().getOpCode()); }
@Test public void testBroadcastChain() { String inputStreamId = "input"; String inputSystem = "input-system"; String inputPhysicalName = "input-stream"; HashMap<String, String> configMap = new HashMap<>(); configMap.put(JobConfig.JOB_NAME, "test-job"); configMap.put(JobConfig.JOB_ID, "1"); StreamTestUtils.addStreamConfigs(configMap, inputStreamId, inputSystem, inputPhysicalName); Config config = new MapConfig(configMap); when(this.context.getJobContext().getConfig()).thenReturn(config); StreamApplicationDescriptorImpl graphSpec = new StreamApplicationDescriptorImpl(appDesc -> { GenericSystemDescriptor sd = new GenericSystemDescriptor(inputSystem, "mockFactoryClass"); GenericInputDescriptor inputDescriptor = sd.getInputDescriptor(inputStreamId, mock(Serde.class)); MessageStream<Object> inputStream = appDesc.getInputStream(inputDescriptor); inputStream.filter(mock(FilterFunction.class)); inputStream.map(mock(MapFunction.class)); }, config); OperatorImplGraph opImplGraph = new OperatorImplGraph(graphSpec.getOperatorSpecGraph(), this.context, mock(Clock.class)); InputOperatorImpl inputOpImpl = opImplGraph.getInputOperator(new SystemStream(inputSystem, inputPhysicalName)); assertEquals(2, inputOpImpl.registeredOperators.size()); assertTrue(inputOpImpl.registeredOperators.stream() .anyMatch(opImpl -> ((OperatorImpl) opImpl).getOperatorSpec().getOpCode() == OpCode.FILTER)); assertTrue(inputOpImpl.registeredOperators.stream() .anyMatch(opImpl -> ((OperatorImpl) opImpl).getOperatorSpec().getOpCode() == OpCode.MAP)); }
@Test public void testJoinChain() { String inputStreamId1 = "input1"; String inputStreamId2 = "input2"; String inputSystem = "input-system"; String inputPhysicalName1 = "input-stream1"; String inputPhysicalName2 = "input-stream2"; HashMap<String, String> configs = new HashMap<>(); configs.put(JobConfig.JOB_NAME, "jobName"); configs.put(JobConfig.JOB_ID, "jobId"); StreamTestUtils.addStreamConfigs(configs, inputStreamId1, inputSystem, inputPhysicalName1); StreamTestUtils.addStreamConfigs(configs, inputStreamId2, inputSystem, inputPhysicalName2); Config config = new MapConfig(configs); when(this.context.getJobContext().getConfig()).thenReturn(config); Integer joinKey = new Integer(1); Function<Object, Integer> keyFn = (Function & Serializable) m -> joinKey; JoinFunction testJoinFunction = new TestJoinFunction("jobName-jobId-join-j1", (BiFunction & Serializable) (m1, m2) -> KV.of(m1, m2), keyFn, keyFn); StreamApplicationDescriptorImpl graphSpec = new StreamApplicationDescriptorImpl(appDesc -> { GenericSystemDescriptor sd = new GenericSystemDescriptor(inputSystem, "mockFactoryClass"); GenericInputDescriptor inputDescriptor1 = sd.getInputDescriptor(inputStreamId1, mock(Serde.class)); GenericInputDescriptor inputDescriptor2 = sd.getInputDescriptor(inputStreamId2, mock(Serde.class)); MessageStream<Object> inputStream1 = appDesc.getInputStream(inputDescriptor1); MessageStream<Object> inputStream2 = appDesc.getInputStream(inputDescriptor2); inputStream1.join(inputStream2, testJoinFunction, mock(Serde.class), mock(Serde.class), mock(Serde.class), Duration.ofHours(1), "j1"); }, config); TaskName mockTaskName = mock(TaskName.class); TaskModel taskModel = mock(TaskModel.class); when(taskModel.getTaskName()).thenReturn(mockTaskName); when(this.context.getTaskContext().getTaskModel()).thenReturn(taskModel); KeyValueStore mockLeftStore = mock(KeyValueStore.class); when(this.context.getTaskContext().getStore(eq("jobName-jobId-join-j1-L"))).thenReturn(mockLeftStore); KeyValueStore mockRightStore = mock(KeyValueStore.class); when(this.context.getTaskContext().getStore(eq("jobName-jobId-join-j1-R"))).thenReturn(mockRightStore); OperatorImplGraph opImplGraph = new OperatorImplGraph(graphSpec.getOperatorSpecGraph(), this.context, mock(Clock.class)); assertEquals(TestJoinFunction.getInstanceByTaskName(mockTaskName, "jobName-jobId-join-j1").numInitCalled, 1); InputOperatorImpl inputOpImpl1 = opImplGraph.getInputOperator(new SystemStream(inputSystem, inputPhysicalName1)); InputOperatorImpl inputOpImpl2 = opImplGraph.getInputOperator(new SystemStream(inputSystem, inputPhysicalName2)); PartialJoinOperatorImpl leftPartialJoinOpImpl = (PartialJoinOperatorImpl) inputOpImpl1.registeredOperators.iterator().next(); PartialJoinOperatorImpl rightPartialJoinOpImpl = (PartialJoinOperatorImpl) inputOpImpl2.registeredOperators.iterator().next(); assertEquals(leftPartialJoinOpImpl.getOperatorSpec(), rightPartialJoinOpImpl.getOperatorSpec()); assertNotSame(leftPartialJoinOpImpl, rightPartialJoinOpImpl); Object mockLeftMessage = mock(Object.class); long currentTimeMillis = System.currentTimeMillis(); when(mockLeftStore.get(eq(joinKey))).thenReturn(new TimestampedValue<>(mockLeftMessage, currentTimeMillis)); IncomingMessageEnvelope leftMessage = new IncomingMessageEnvelope(mock(SystemStreamPartition.class), "", "", mockLeftMessage); inputOpImpl1.onMessage(leftMessage, mock(MessageCollector.class), mock(TaskCoordinator.class)); Object mockRightMessage = mock(Object.class); when(mockRightStore.get(eq(joinKey))).thenReturn(new TimestampedValue<>(mockRightMessage, currentTimeMillis)); IncomingMessageEnvelope rightMessage = new IncomingMessageEnvelope(mock(SystemStreamPartition.class), "", "", mockRightMessage); inputOpImpl2.onMessage(rightMessage, mock(MessageCollector.class), mock(TaskCoordinator.class)); assertEquals(((TestJoinFunction) TestJoinFunction.getInstanceByTaskName(mockTaskName, "jobName-jobId-join-j1")).joinResults.size(), 1); KV joinResult = (KV) ((TestJoinFunction) TestJoinFunction.getInstanceByTaskName(mockTaskName, "jobName-jobId-join-j1")).joinResults.iterator().next(); assertEquals(joinResult.getKey(), mockLeftMessage); assertEquals(joinResult.getValue(), mockRightMessage); }