Java Code Examples for org.apache.samza.system.descriptors.GenericSystemDescriptor#getOutputDescriptor()

@Test
public void testGetOutputStreamWithKeyValueSerde() {
  String streamId = "test-stream-1";
  KVSerde mockKVSerde = mock(KVSerde.class);
  Serde mockKeySerde = mock(Serde.class);
  Serde mockValueSerde = mock(Serde.class);
  doReturn(mockKeySerde).when(mockKVSerde).getKeySerde();
  doReturn(mockValueSerde).when(mockKVSerde).getValueSerde();
  GenericSystemDescriptor sd = new GenericSystemDescriptor("mockSystem", "mockSystemFactoryClass");
  GenericOutputDescriptor osd = sd.getOutputDescriptor(streamId, mockKVSerde);

  StreamApplicationDescriptorImpl streamAppDesc = new StreamApplicationDescriptorImpl(appDesc -> {
    appDesc.getOutputStream(osd);
  }, getConfig());

  OutputStreamImpl<TestMessageEnvelope> outputStreamImpl = streamAppDesc.getOutputStreams().get(streamId);
  assertEquals(streamId, outputStreamImpl.getStreamId());
  assertEquals(osd, streamAppDesc.getOutputDescriptors().get(streamId));
  assertEquals(mockKeySerde, outputStreamImpl.getKeySerde());
  assertEquals(mockValueSerde, outputStreamImpl.getValueSerde());
}
 

@Test
public void testGetOutputStreamWithValueSerde() {
  String streamId = "test-stream-1";
  Serde mockValueSerde = mock(Serde.class);
  GenericSystemDescriptor sd = new GenericSystemDescriptor("mockSystem", "mockSystemFactoryClass");
  GenericOutputDescriptor osd = sd.getOutputDescriptor(streamId, mockValueSerde);

  StreamApplicationDescriptorImpl streamAppDesc = new StreamApplicationDescriptorImpl(appDesc -> {
    appDesc.getOutputStream(osd);
  }, getConfig());

  OutputStreamImpl<TestMessageEnvelope> outputStreamImpl = streamAppDesc.getOutputStreams().get(streamId);
  assertEquals(streamId, outputStreamImpl.getStreamId());
  assertEquals(osd, streamAppDesc.getOutputDescriptors().get(streamId));
  assertTrue(outputStreamImpl.getKeySerde() instanceof NoOpSerde);
  assertEquals(mockValueSerde, outputStreamImpl.getValueSerde());
}
 

@Test(expected = IllegalArgumentException.class)
public void testGetOutputStreamWithNullSerde() {
  String streamId = "test-stream-1";
  GenericSystemDescriptor sd = new GenericSystemDescriptor("mockSystem", "mockSystemFactoryClass");
  GenericOutputDescriptor osd = sd.getOutputDescriptor(streamId, null);
  new StreamApplicationDescriptorImpl(appDesc -> {
    appDesc.getOutputStream(osd);
  }, getConfig());
}
 

@Test(expected = IllegalStateException.class)
public void testSetDefaultSystemDescriptorAfterGettingOutputStream() {
  String streamId = "test-stream-1";
  GenericSystemDescriptor sd = new GenericSystemDescriptor("mockSystem", "mockSystemFactoryClass");
  GenericOutputDescriptor osd = sd.getOutputDescriptor(streamId, mock(Serde.class));
  new StreamApplicationDescriptorImpl(appDesc -> {
    appDesc.getOutputStream(osd);
    appDesc.withDefaultSystem(sd); // should throw exception
  }, getConfig());
}
 

@Test(expected = IllegalStateException.class)
public void testGetSameOutputStreamTwice() {
  String streamId = "test-stream-1";
  GenericSystemDescriptor sd = new GenericSystemDescriptor("mockSystem", "mockSystemFactoryClass");
  GenericOutputDescriptor osd1 = sd.getOutputDescriptor(streamId, mock(Serde.class));
  GenericOutputDescriptor osd2 = sd.getOutputDescriptor(streamId, mock(Serde.class));
  new StreamApplicationDescriptorImpl(appDesc -> {
    appDesc.getOutputStream(osd1);
    appDesc.getOutputStream(osd2); // should throw exception
  }, getConfig());
}
 

@Before
public void setUp() {
  defaultSerde = KVSerde.of(new StringSerde(), new JsonSerdeV2<>());
  inputSystemDescriptor = new GenericSystemDescriptor("input-system", "mockSystemFactoryClassName");
  outputSystemDescriptor = new GenericSystemDescriptor("output-system", "mockSystemFactoryClassName");
  intermediateSystemDescriptor = new GenericSystemDescriptor("intermediate-system", "mockSystemFactoryClassName");
  input1Descriptor = inputSystemDescriptor.getInputDescriptor("input1", defaultSerde);
  input2Descriptor = inputSystemDescriptor.getInputDescriptor("input2", defaultSerde);
  outputDescriptor = outputSystemDescriptor.getOutputDescriptor("output", defaultSerde);
  intermediateInputDescriptor = intermediateSystemDescriptor.getInputDescriptor("jobName-jobId-partition_by-p1", defaultSerde)
      .withPhysicalName("jobName-jobId-partition_by-p1");
  intermediateOutputDescriptor = intermediateSystemDescriptor.getOutputDescriptor("jobName-jobId-partition_by-p1", defaultSerde)
      .withPhysicalName("jobName-jobId-partition_by-p1");
  broadcastInputDesriptor = intermediateSystemDescriptor.getInputDescriptor("jobName-jobId-broadcast-b1", defaultSerde)
      .withPhysicalName("jobName-jobId-broadcast-b1");

  Map<String, String> configs = new HashMap<>();
  configs.put(JobConfig.JOB_NAME, "jobName");
  configs.put(JobConfig.JOB_ID, "jobId");
  configs.putAll(input1Descriptor.toConfig());
  configs.putAll(input2Descriptor.toConfig());
  configs.putAll(outputDescriptor.toConfig());
  configs.putAll(inputSystemDescriptor.toConfig());
  configs.putAll(outputSystemDescriptor.toConfig());
  configs.putAll(intermediateSystemDescriptor.toConfig());
  configs.put(JobConfig.JOB_DEFAULT_SYSTEM, intermediateSystemDescriptor.getSystemName());
  mockConfig = spy(new MapConfig(configs));

  mockStreamAppDesc = new StreamApplicationDescriptorImpl(getRepartitionJoinStreamApplication(), mockConfig);
}
 

@Override
public void describe(StreamApplicationDescriptor appDescriptor) {
  // Define a system descriptor for Kafka
  KafkaSystemDescriptor kafkaSystemDescriptor =
      new KafkaSystemDescriptor("kafka").withConsumerZkConnect(KAFKA_CONSUMER_ZK_CONNECT)
          .withProducerBootstrapServers(KAFKA_PRODUCER_BOOTSTRAP_SERVERS)
          .withDefaultStreamConfigs(KAFKA_DEFAULT_STREAM_CONFIGS);

  KafkaInputDescriptor<PageView> pageViewInputDescriptor =
      kafkaSystemDescriptor.getInputDescriptor(INPUT_PAGEVIEW_STREAM_ID, new JsonSerdeV2<>(PageView.class));

  // Define a system descriptor for Azure Blob Storage
  GenericSystemDescriptor azureBlobSystemDescriptor =
      new GenericSystemDescriptor(OUTPUT_SYSTEM, "org.apache.samza.system.azureblob.AzureBlobSystemFactory");

  GenericOutputDescriptor<PageViewAvroRecord> azureBlobOuputDescriptor =
      azureBlobSystemDescriptor.getOutputDescriptor(OUTPUT_STREAM, new NoOpSerde<>());

  // Set Kafka as the default system for the job
  appDescriptor.withDefaultSystem(kafkaSystemDescriptor);

  // Define the input and output streams with descriptors
  MessageStream<PageView> pageViewInput = appDescriptor.getInputStream(pageViewInputDescriptor);
  OutputStream<PageViewAvroRecord> pageViewAvroRecordOutputStream = appDescriptor.getOutputStream(azureBlobOuputDescriptor);

  // Define the execution flow with the high-level API
  pageViewInput
      .map((message) -> {
        LOG.info("Sending: Received PageViewEvent with pageId: " + message.pageId);
        return PageViewAvroRecord.buildPageViewRecord(message);
      })
      .sendTo(pageViewAvroRecordOutputStream);
}
 

@Before
public void setUp() {
  input1Spec = new StreamSpec("input1", "input1", "input-system");
  input2Spec = new StreamSpec("input2", "input2", "input-system");
  outputSpec = new StreamSpec("output", "output", "output-system");
  repartitionSpec =
      new StreamSpec("jobName-jobId-partition_by-p1", "partition_by-p1", "intermediate-system");


  defaultSerde = KVSerde.of(new StringSerde(), new JsonSerdeV2<>());
  inputSystemDescriptor = new GenericSystemDescriptor("input-system", "mockSystemFactoryClassName");
  outputSystemDescriptor = new GenericSystemDescriptor("output-system", "mockSystemFactoryClassName");
  intermediateSystemDescriptor = new GenericSystemDescriptor("intermediate-system", "mockSystemFactoryClassName");
  input1Descriptor = inputSystemDescriptor.getInputDescriptor("input1", defaultSerde);
  input2Descriptor = inputSystemDescriptor.getInputDescriptor("input2", defaultSerde);
  outputDescriptor = outputSystemDescriptor.getOutputDescriptor("output", defaultSerde);
  table1Descriptor = new TestLocalTableDescriptor.MockLocalTableDescriptor("table1", defaultSerde);
  table2Descriptor = new TestLocalTableDescriptor.MockLocalTableDescriptor("table2", defaultSerde);

  Map<String, String> configs = new HashMap<>();
  configs.put(JobConfig.JOB_NAME, "jobName");
  configs.put(JobConfig.JOB_ID, "jobId");
  mockConfig = spy(new MapConfig(configs));

  mockJobNode = mock(JobNode.class);
  StreamEdge input1Edge = new StreamEdge(input1Spec, false, false, mockConfig);
  StreamEdge input2Edge = new StreamEdge(input2Spec, false, false, mockConfig);
  StreamEdge outputEdge = new StreamEdge(outputSpec, false, false, mockConfig);
  StreamEdge repartitionEdge = new StreamEdge(repartitionSpec, true, false, mockConfig);
  Map<String, StreamEdge> inputEdges = new HashMap<>();
  inputEdges.put(input1Descriptor.getStreamId(), input1Edge);
  inputEdges.put(input2Descriptor.getStreamId(), input2Edge);
  inputEdges.put(repartitionSpec.getId(), repartitionEdge);
  Map<String, StreamEdge> outputEdges = new HashMap<>();
  outputEdges.put(outputDescriptor.getStreamId(), outputEdge);
  outputEdges.put(repartitionSpec.getId(), repartitionEdge);
  when(mockJobNode.getInEdges()).thenReturn(inputEdges);
  when(mockJobNode.getOutEdges()).thenReturn(outputEdges);
  when(mockJobNode.getConfig()).thenReturn(mockConfig);
  when(mockJobNode.getJobName()).thenReturn("jobName");
  when(mockJobNode.getJobId()).thenReturn("jobId");
  when(mockJobNode.getJobNameAndId()).thenReturn(JobNode.createJobNameAndId("jobName", "jobId"));

  Map<String, TableDescriptor> tables = new HashMap<>();
  tables.put(table1Descriptor.getTableId(), table1Descriptor);
  tables.put(table2Descriptor.getTableId(), table2Descriptor);
  when(mockJobNode.getTables()).thenReturn(tables);
}
 

@Test
public void testRepartitionedJoinStreamApplication() throws Exception {

  /**
   * the graph looks like the following.
   * number in parentheses () indicates number of stream partitions.
   * number in parentheses in quotes ("") indicates expected partition count.
   * number in square brackets [] indicates operator ID.
   *
   * input3 (32) -> filter [7] -> partitionBy [8] ("64") -> map [10] -> join [14] -> sendTo(output2) [15] (16)
   *                                                                   |
   *              input2 (16) -> partitionBy [3] ("64") -> filter [5] -| -> sink [13]
   *                                                                   |
   *                                         input1 (64) -> map [1] -> join [11] -> sendTo(output1) [12] (8)
   *
   */

  Map<String, String> configMap = new HashMap<>();
  configMap.put(JobConfig.JOB_NAME, "test-app");
  configMap.put(JobConfig.JOB_DEFAULT_SYSTEM, "test-system");
  StreamTestUtils.addStreamConfigs(configMap, "input1", "system1", "input1");
  StreamTestUtils.addStreamConfigs(configMap, "input2", "system2", "input2");
  StreamTestUtils.addStreamConfigs(configMap, "input3", "system2", "input3");
  StreamTestUtils.addStreamConfigs(configMap, "output1", "system1", "output1");
  StreamTestUtils.addStreamConfigs(configMap, "output2", "system2", "output2");
  Config config = new MapConfig(configMap);

  // set up external partition count
  Map<String, Integer> system1Map = new HashMap<>();
  system1Map.put("input1", 64);
  system1Map.put("output1", 8);
  Map<String, Integer> system2Map = new HashMap<>();
  system2Map.put("input2", 16);
  system2Map.put("input3", 32);
  system2Map.put("output2", 16);

  SystemAdmin systemAdmin1 = createSystemAdmin(system1Map);
  SystemAdmin systemAdmin2 = createSystemAdmin(system2Map);
  SystemAdmins systemAdmins = mock(SystemAdmins.class);
  when(systemAdmins.getSystemAdmin("system1")).thenReturn(systemAdmin1);
  when(systemAdmins.getSystemAdmin("system2")).thenReturn(systemAdmin2);
  StreamManager streamManager = new StreamManager(systemAdmins);

  StreamApplicationDescriptorImpl graphSpec = new StreamApplicationDescriptorImpl(appDesc -> {
    KVSerde<Object, Object> kvSerde = new KVSerde<>(new NoOpSerde(), new NoOpSerde());
    String mockSystemFactoryClass = "factory.class.name";
    GenericSystemDescriptor system1 = new GenericSystemDescriptor("system1", mockSystemFactoryClass);
    GenericSystemDescriptor system2 = new GenericSystemDescriptor("system2", mockSystemFactoryClass);
    GenericInputDescriptor<KV<Object, Object>> input1Descriptor = system1.getInputDescriptor("input1", kvSerde);
    GenericInputDescriptor<KV<Object, Object>> input2Descriptor = system2.getInputDescriptor("input2", kvSerde);
    GenericInputDescriptor<KV<Object, Object>> input3Descriptor = system2.getInputDescriptor("input3", kvSerde);
    GenericOutputDescriptor<KV<Object, Object>>  output1Descriptor = system1.getOutputDescriptor("output1", kvSerde);
    GenericOutputDescriptor<KV<Object, Object>> output2Descriptor = system2.getOutputDescriptor("output2", kvSerde);

    MessageStream<KV<Object, Object>> messageStream1 =
        appDesc.getInputStream(input1Descriptor)
            .map(m -> m);
    MessageStream<KV<Object, Object>> messageStream2 =
        appDesc.getInputStream(input2Descriptor)
            .partitionBy(m -> m.key, m -> m.value, mock(KVSerde.class), "p1")
            .filter(m -> true);
    MessageStream<KV<Object, Object>> messageStream3 =
        appDesc.getInputStream(input3Descriptor)
            .filter(m -> true)
            .partitionBy(m -> m.key, m -> m.value, mock(KVSerde.class), "p2")
            .map(m -> m);
    OutputStream<KV<Object, Object>> outputStream1 = appDesc.getOutputStream(output1Descriptor);
    OutputStream<KV<Object, Object>> outputStream2 = appDesc.getOutputStream(output2Descriptor);

    messageStream1
        .join(messageStream2,
            (JoinFunction<Object, KV<Object, Object>, KV<Object, Object>, KV<Object, Object>>) mock(JoinFunction.class),
            mock(Serde.class), mock(Serde.class), mock(Serde.class), Duration.ofHours(2), "j1")
        .sendTo(outputStream1);
    messageStream2.sink((message, collector, coordinator) -> { });
    messageStream3
        .join(messageStream2,
            (JoinFunction<Object, KV<Object, Object>, KV<Object, Object>, KV<Object, Object>>) mock(JoinFunction.class),
            mock(Serde.class), mock(Serde.class), mock(Serde.class), Duration.ofHours(1), "j2")
        .sendTo(outputStream2);
  }, config);

  ExecutionPlanner planner = new ExecutionPlanner(config, streamManager);
  ExecutionPlan plan = planner.plan(graphSpec);
  String json = plan.getPlanAsJson();
  System.out.println(json);

  // deserialize
  ObjectMapper mapper = new ObjectMapper();
  JobGraphJsonGenerator.JobGraphJson nodes = mapper.readValue(json, JobGraphJsonGenerator.JobGraphJson.class);
  assertEquals(5, nodes.jobs.get(0).operatorGraph.inputStreams.size());
  assertEquals(11, nodes.jobs.get(0).operatorGraph.operators.size());
  assertEquals(3, nodes.sourceStreams.size());
  assertEquals(2, nodes.sinkStreams.size());
  assertEquals(2, nodes.intermediateStreams.size());
}
 

@Test
public void testRepartitionedWindowStreamApplication() throws Exception {
  Map<String, String> configMap = new HashMap<>();
  configMap.put(JobConfig.JOB_NAME, "test-app");
  configMap.put(JobConfig.JOB_DEFAULT_SYSTEM, "test-system");
  StreamTestUtils.addStreamConfigs(configMap, "PageView", "hdfs", "hdfs:/user/dummy/PageViewEvent");
  StreamTestUtils.addStreamConfigs(configMap, "PageViewCount", "kafka", "PageViewCount");
  Config config = new MapConfig(configMap);

  // set up external partition count
  Map<String, Integer> system1Map = new HashMap<>();
  system1Map.put("hdfs:/user/dummy/PageViewEvent", 512);
  Map<String, Integer> system2Map = new HashMap<>();
  system2Map.put("PageViewCount", 16);

  SystemAdmin systemAdmin1 = createSystemAdmin(system1Map);
  SystemAdmin systemAdmin2 = createSystemAdmin(system2Map);
  SystemAdmins systemAdmins = mock(SystemAdmins.class);
  when(systemAdmins.getSystemAdmin("hdfs")).thenReturn(systemAdmin1);
  when(systemAdmins.getSystemAdmin("kafka")).thenReturn(systemAdmin2);
  StreamManager streamManager = new StreamManager(systemAdmins);

  StreamApplicationDescriptorImpl graphSpec = new StreamApplicationDescriptorImpl(appDesc -> {
    KVSerde<String, PageViewEvent> pvSerde = KVSerde.of(new StringSerde(), new JsonSerdeV2<>(PageViewEvent.class));
    GenericSystemDescriptor isd = new GenericSystemDescriptor("hdfs", "mockSystemFactoryClass");
    GenericInputDescriptor<KV<String, PageViewEvent>> pageView = isd.getInputDescriptor("PageView", pvSerde);

    KVSerde<String, Long> pvcSerde = KVSerde.of(new StringSerde(), new LongSerde());
    GenericSystemDescriptor osd = new GenericSystemDescriptor("kafka", "mockSystemFactoryClass");
    GenericOutputDescriptor<KV<String, Long>> pageViewCount = osd.getOutputDescriptor("PageViewCount", pvcSerde);

    MessageStream<KV<String, PageViewEvent>> inputStream = appDesc.getInputStream(pageView);
    OutputStream<KV<String, Long>> outputStream = appDesc.getOutputStream(pageViewCount);
    inputStream
        .partitionBy(kv -> kv.getValue().getCountry(), kv -> kv.getValue(), pvSerde, "keyed-by-country")
        .window(Windows.keyedTumblingWindow(kv -> kv.getValue().getCountry(),
            Duration.ofSeconds(10L), () -> 0L, (m, c) -> c + 1L, new StringSerde(), new LongSerde()),
            "count-by-country")
        .map(pane -> new KV<>(pane.getKey().getKey(), pane.getMessage()))
        .sendTo(outputStream);
  }, config);

  ExecutionPlanner planner = new ExecutionPlanner(config, streamManager);
  ExecutionPlan plan = planner.plan(graphSpec);
  String json = plan.getPlanAsJson();
  System.out.println(json);

  // deserialize
  ObjectMapper mapper = new ObjectMapper();
  JobGraphJsonGenerator.JobGraphJson nodes = mapper.readValue(json, JobGraphJsonGenerator.JobGraphJson.class);
  JobGraphJsonGenerator.OperatorGraphJson operatorGraphJson = nodes.jobs.get(0).operatorGraph;
  assertEquals(2, operatorGraphJson.inputStreams.size());
  assertEquals(4, operatorGraphJson.operators.size());
  assertEquals(1, nodes.sourceStreams.size());
  assertEquals(1, nodes.sinkStreams.size());
  assertEquals(1, nodes.intermediateStreams.size());

  // verify partitionBy op output to the intermdiate stream of the same id
  assertEquals(operatorGraphJson.operators.get("test-app-1-partition_by-keyed-by-country").get("outputStreamId"),
      "test-app-1-partition_by-keyed-by-country");
  assertEquals(operatorGraphJson.operators.get("test-app-1-send_to-5").get("outputStreamId"),
      "PageViewCount");
}
 

@Before
public void setup() {
  Map<String, String> configMap = new HashMap<>();
  configMap.put(JobConfig.JOB_NAME, "test-app");
  configMap.put(JobConfig.JOB_DEFAULT_SYSTEM, DEFAULT_SYSTEM);
  StreamTestUtils.addStreamConfigs(configMap, "input1", "system1", "input1");
  StreamTestUtils.addStreamConfigs(configMap, "input2", "system2", "input2");
  StreamTestUtils.addStreamConfigs(configMap, "input3", "system2", "input3");
  StreamTestUtils.addStreamConfigs(configMap, "input4", "system1", "input4");
  StreamTestUtils.addStreamConfigs(configMap, "output1", "system1", "output1");
  StreamTestUtils.addStreamConfigs(configMap, "output2", "system2", "output2");
  config = new MapConfig(configMap);

  input1Spec = new StreamSpec("input1", "input1", "system1");
  input2Spec = new StreamSpec("input2", "input2", "system2");
  input3Spec = new StreamSpec("input3", "input3", "system2");

  output1Spec = new StreamSpec("output1", "output1", "system1");
  output2Spec = new StreamSpec("output2", "output2", "system2");

  KVSerde<Object, Object> kvSerde = new KVSerde<>(new NoOpSerde(), new NoOpSerde());
  String mockSystemFactoryClass = "factory.class.name";
  system1Descriptor = new GenericSystemDescriptor("system1", mockSystemFactoryClass);
  system2Descriptor = new GenericSystemDescriptor("system2", mockSystemFactoryClass);
  input1Descriptor = system1Descriptor.getInputDescriptor("input1", kvSerde);
  input2Descriptor = system2Descriptor.getInputDescriptor("input2", kvSerde);
  input3Descriptor = system2Descriptor.getInputDescriptor("input3", kvSerde);
  input4Descriptor = system1Descriptor.getInputDescriptor("input4", kvSerde);
  output1Descriptor = system1Descriptor.getOutputDescriptor("output1", kvSerde);
  output2Descriptor = system2Descriptor.getOutputDescriptor("output2", kvSerde);

  // clean and set up sets and maps of descriptors
  systemDescriptors.clear();
  inputDescriptors.clear();
  outputDescriptors.clear();
  tableDescriptors.clear();
  systemDescriptors.add(system1Descriptor);
  systemDescriptors.add(system2Descriptor);
  inputDescriptors.put(input1Descriptor.getStreamId(), input1Descriptor);
  inputDescriptors.put(input2Descriptor.getStreamId(), input2Descriptor);
  inputDescriptors.put(input3Descriptor.getStreamId(), input3Descriptor);
  inputDescriptors.put(input4Descriptor.getStreamId(), input4Descriptor);
  outputDescriptors.put(output1Descriptor.getStreamId(), output1Descriptor);
  outputDescriptors.put(output2Descriptor.getStreamId(), output2Descriptor);


  // set up external partition count
  Map<String, Integer> system1Map = new HashMap<>();
  system1Map.put("input1", 64);
  system1Map.put("output1", 8);
  system1Map.put("input4", IntermediateStreamManager.MAX_INFERRED_PARTITIONS * 2);
  Map<String, Integer> system2Map = new HashMap<>();
  system2Map.put("input2", 16);
  system2Map.put("input3", 32);
  system2Map.put("output2", 16);

  SystemAdmin systemAdmin1 = createSystemAdmin(system1Map);
  SystemAdmin systemAdmin2 = createSystemAdmin(system2Map);
  systemAdmins = mock(SystemAdmins.class);
  when(systemAdmins.getSystemAdmin("system1")).thenReturn(systemAdmin1);
  when(systemAdmins.getSystemAdmin("system2")).thenReturn(systemAdmin2);
  streamManager = new StreamManager(systemAdmins);
}
 

@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()); // is terminal but paired with an input operator
  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 testGetOutputToInputStreams() {
  String inputStreamId1 = "input1";
  String inputStreamId2 = "input2";
  String inputStreamId3 = "input3";
  String inputSystem = "input-system";

  String outputStreamId1 = "output1";
  String outputStreamId2 = "output2";
  String outputSystem = "output-system";

  String intStreamId1 = "test-app-1-partition_by-p1";
  String intStreamId2 = "test-app-1-partition_by-p2";
  String intSystem = "test-system";

  HashMap<String, String> configs = new HashMap<>();
  configs.put(JobConfig.JOB_NAME, "test-app");
  configs.put(JobConfig.JOB_DEFAULT_SYSTEM, intSystem);
  StreamTestUtils.addStreamConfigs(configs, inputStreamId1, inputSystem, inputStreamId1);
  StreamTestUtils.addStreamConfigs(configs, inputStreamId2, inputSystem, inputStreamId2);
  StreamTestUtils.addStreamConfigs(configs, inputStreamId3, inputSystem, inputStreamId3);
  StreamTestUtils.addStreamConfigs(configs, outputStreamId1, outputSystem, outputStreamId1);
  StreamTestUtils.addStreamConfigs(configs, outputStreamId2, outputSystem, outputStreamId2);
  Config config = new MapConfig(configs);
  when(this.context.getJobContext().getConfig()).thenReturn(config);

  StreamApplicationDescriptorImpl graphSpec = new StreamApplicationDescriptorImpl(appDesc -> {
    GenericSystemDescriptor isd = new GenericSystemDescriptor(inputSystem, "mockFactoryClass");
    GenericInputDescriptor inputDescriptor1 = isd.getInputDescriptor(inputStreamId1, mock(Serde.class));
    GenericInputDescriptor inputDescriptor2 = isd.getInputDescriptor(inputStreamId2, mock(Serde.class));
    GenericInputDescriptor inputDescriptor3 = isd.getInputDescriptor(inputStreamId3, mock(Serde.class));
    GenericSystemDescriptor osd = new GenericSystemDescriptor(outputSystem, "mockFactoryClass");
    GenericOutputDescriptor outputDescriptor1 = osd.getOutputDescriptor(outputStreamId1, mock(Serde.class));
    GenericOutputDescriptor outputDescriptor2 = osd.getOutputDescriptor(outputStreamId2, mock(Serde.class));
    MessageStream messageStream1 = appDesc.getInputStream(inputDescriptor1).map(m -> m);
    MessageStream messageStream2 = appDesc.getInputStream(inputDescriptor2).filter(m -> true);
    MessageStream messageStream3 =
        appDesc.getInputStream(inputDescriptor3)
            .filter(m -> true)
            .partitionBy(m -> "m", m -> m, mock(KVSerde.class),  "p1")
            .map(m -> m);
    OutputStream<Object> outputStream1 = appDesc.getOutputStream(outputDescriptor1);
    OutputStream<Object> outputStream2 = appDesc.getOutputStream(outputDescriptor2);

    messageStream1
        .join(messageStream2, mock(JoinFunction.class),
            mock(Serde.class), mock(Serde.class), mock(Serde.class), Duration.ofHours(2), "j1")
        .partitionBy(m -> "m", m -> m, mock(KVSerde.class), "p2")
        .sendTo(outputStream1);
    messageStream3
        .join(messageStream2, mock(JoinFunction.class),
            mock(Serde.class), mock(Serde.class), mock(Serde.class), Duration.ofHours(1), "j2")
        .sendTo(outputStream2);
  }, config);

  Multimap<SystemStream, SystemStream> outputToInput =
      OperatorImplGraph.getIntermediateToInputStreamsMap(graphSpec.getOperatorSpecGraph(), new StreamConfig(config));
  Collection<SystemStream> inputs = outputToInput.get(new SystemStream(intSystem, intStreamId2));
  assertEquals(inputs.size(), 2);
  assertTrue(inputs.contains(new SystemStream(inputSystem, inputStreamId1)));
  assertTrue(inputs.contains(new SystemStream(inputSystem, inputStreamId2)));

  inputs = outputToInput.get(new SystemStream(intSystem, intStreamId1));
  assertEquals(inputs.size(), 1);
  assertEquals(inputs.iterator().next(), new SystemStream(inputSystem, inputStreamId3));
}