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

@Override
public void translatePortable(
    PipelineNode.PTransformNode transform,
    QueryablePipeline pipeline,
    PortableTranslationContext ctx) {

  final String outputId = ctx.getOutputId(transform);
  final GenericSystemDescriptor systemDescriptor =
      new GenericSystemDescriptor(outputId, SamzaImpulseSystemFactory.class.getName());

  // The KvCoder is needed here for Samza not to crop the key.
  final Serde<KV<?, OpMessage<byte[]>>> kvSerde = KVSerde.of(new NoOpSerde(), new NoOpSerde<>());
  final GenericInputDescriptor<KV<?, OpMessage<byte[]>>> inputDescriptor =
      systemDescriptor.getInputDescriptor(outputId, kvSerde);

  ctx.registerInputMessageStream(outputId, inputDescriptor);
}
 

private StreamApplicationDescriptorImpl getKeyedTumblingWindowStreamGraph(AccumulationMode mode,
    Duration duration, Trigger<KV<Integer, Integer>> earlyTrigger) throws IOException {

  StreamApplication userApp = appDesc -> {
    KVSerde<Integer, Integer> kvSerde = KVSerde.of(new IntegerSerde(), new IntegerSerde());
    GenericSystemDescriptor sd = new GenericSystemDescriptor("kafka", "mockFactoryClass");
    GenericInputDescriptor<KV<Integer, Integer>> inputDescriptor = sd.getInputDescriptor("integers", kvSerde);
    appDesc.getInputStream(inputDescriptor)
        .window(Windows.keyedTumblingWindow(KV::getKey, duration, new IntegerSerde(), kvSerde)
            .setEarlyTrigger(earlyTrigger).setAccumulationMode(mode), "w1")
        .sink((message, messageCollector, taskCoordinator) -> {
          SystemStream outputSystemStream = new SystemStream("outputSystem", "outputStream");
          messageCollector.send(new OutgoingMessageEnvelope(outputSystemStream, message));
        });
  };

  return new StreamApplicationDescriptorImpl(userApp, config);
}
 

private StreamApplicationDescriptorImpl getKeyedSessionWindowStreamGraph(AccumulationMode mode, Duration duration) throws IOException {
  StreamApplication userApp = appDesc -> {
    KVSerde<Integer, Integer> kvSerde = KVSerde.of(new IntegerSerde(), new IntegerSerde());
    GenericSystemDescriptor sd = new GenericSystemDescriptor("kafka", "mockFactoryClass");
    GenericInputDescriptor<KV<Integer, Integer>> inputDescriptor = sd.getInputDescriptor("integers", kvSerde);
    appDesc.getInputStream(inputDescriptor)
        .window(Windows.keyedSessionWindow(KV::getKey, duration, new IntegerSerde(), kvSerde)
            .setAccumulationMode(mode), "w1")
        .sink((message, messageCollector, taskCoordinator) -> {
          SystemStream outputSystemStream = new SystemStream("outputSystem", "outputStream");
          messageCollector.send(new OutgoingMessageEnvelope(outputSystemStream, message));
        });
  };

  return new StreamApplicationDescriptorImpl(userApp, config);
}
 

@Test
public void testGetInputStreamWithKeyValueSerde() {

  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");
  GenericInputDescriptor isd = sd.getInputDescriptor(streamId, mockKVSerde);
  StreamApplicationDescriptorImpl streamAppDesc = new StreamApplicationDescriptorImpl(appDesc -> {
    appDesc.getInputStream(isd);
  }, getConfig());

  InputOperatorSpec inputOpSpec = streamAppDesc.getInputOperators().get(streamId);
  assertEquals(OpCode.INPUT, inputOpSpec.getOpCode());
  assertEquals(streamId, inputOpSpec.getStreamId());
  assertEquals(isd, streamAppDesc.getInputDescriptors().get(streamId));
  assertEquals(mockKeySerde, inputOpSpec.getKeySerde());
  assertEquals(mockValueSerde, inputOpSpec.getValueSerde());
}
 

/** The dummy stream created will only be used in Beam tests. */
private static InputDescriptor<OpMessage<String>, ?> createDummyStreamDescriptor(String id) {
  final GenericSystemDescriptor dummySystem =
      new GenericSystemDescriptor(id, InMemorySystemFactory.class.getName());
  final GenericInputDescriptor<OpMessage<String>> dummyInput =
      dummySystem.getInputDescriptor(id, new NoOpSerde<>());
  dummyInput.withOffsetDefault(SystemStreamMetadata.OffsetType.OLDEST);
  final Config config = new MapConfig(dummyInput.toConfig(), dummySystem.toConfig());
  final SystemFactory factory = new InMemorySystemFactory();
  final StreamSpec dummyStreamSpec = new StreamSpec(id, id, id, 1);
  factory.getAdmin(id, config).createStream(dummyStreamSpec);

  final SystemProducer producer = factory.getProducer(id, config, null);
  final SystemStream sysStream = new SystemStream(id, id);
  final Consumer<Object> sendFn =
      (msg) -> {
        producer.send(id, new OutgoingMessageEnvelope(sysStream, 0, null, msg));
      };
  final WindowedValue<String> windowedValue =
      WindowedValue.timestampedValueInGlobalWindow("dummy", new Instant());

  sendFn.accept(OpMessage.ofElement(windowedValue));
  sendFn.accept(new WatermarkMessage(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis()));
  sendFn.accept(new EndOfStreamMessage(null));
  return dummyInput;
}
 

private StreamApplicationDescriptorImpl getTumblingWindowStreamGraph(AccumulationMode mode,
    Duration duration, Trigger<KV<Integer, Integer>> earlyTrigger) throws IOException {
  StreamApplication userApp = appDesc -> {
    KVSerde<Integer, Integer> kvSerde = KVSerde.of(new IntegerSerde(), new IntegerSerde());
    GenericSystemDescriptor sd = new GenericSystemDescriptor("kafka", "mockFactoryClass");
    GenericInputDescriptor<KV<Integer, Integer>> inputDescriptor = sd.getInputDescriptor("integers", kvSerde);
    appDesc.getInputStream(inputDescriptor)
        .window(Windows.tumblingWindow(duration, kvSerde).setEarlyTrigger(earlyTrigger)
            .setAccumulationMode(mode), "w1")
        .sink((message, messageCollector, taskCoordinator) -> {
          SystemStream outputSystemStream = new SystemStream("outputSystem", "outputStream");
          messageCollector.send(new OutgoingMessageEnvelope(outputSystemStream, message));
        });
  };

  return new StreamApplicationDescriptorImpl(userApp, config);
}
 

@Test
public void testMultipleGetInputStreams() {
  String streamId1 = "test-stream-1";
  String streamId2 = "test-stream-2";
  GenericSystemDescriptor sd = new GenericSystemDescriptor("mockSystem", "mockSystemFactoryClass");
  GenericInputDescriptor isd1 = sd.getInputDescriptor(streamId1, mock(Serde.class));
  GenericInputDescriptor isd2 = sd.getInputDescriptor(streamId2, mock(Serde.class));

  StreamApplicationDescriptorImpl streamAppDesc = new StreamApplicationDescriptorImpl(appDesc -> {
    appDesc.getInputStream(isd1);
    appDesc.getInputStream(isd2);
  }, getConfig());

  InputOperatorSpec inputOpSpec1 = streamAppDesc.getInputOperators().get(streamId1);
  InputOperatorSpec inputOpSpec2 = streamAppDesc.getInputOperators().get(streamId2);

  assertEquals(2, streamAppDesc.getInputOperators().size());
  assertEquals(streamId1, inputOpSpec1.getStreamId());
  assertEquals(streamId2, inputOpSpec2.getStreamId());
  assertEquals(2, streamAppDesc.getInputDescriptors().size());
  assertEquals(isd1, streamAppDesc.getInputDescriptors().get(streamId1));
  assertEquals(isd2, streamAppDesc.getInputDescriptors().get(streamId2));
}
 

@Test(expected = SamzaException.class)
public void joinWithSelfThrowsException() throws Exception {
  Map<String, String> mapConfig = new HashMap<>();
  mapConfig.put("job.name", "jobName");
  mapConfig.put("job.id", "jobId");
  StreamTestUtils.addStreamConfigs(mapConfig, "inStream", "insystem", "instream");
  Config config = new MapConfig(mapConfig);

  StreamApplicationDescriptorImpl streamAppDesc = new StreamApplicationDescriptorImpl(appDesc -> {
    IntegerSerde integerSerde = new IntegerSerde();
    KVSerde<Integer, Integer> kvSerde = KVSerde.of(integerSerde, integerSerde);
    GenericSystemDescriptor sd = new GenericSystemDescriptor("insystem", "mockFactoryClassName");
    GenericInputDescriptor<KV<Integer, Integer>> inputDescriptor = sd.getInputDescriptor("inStream", kvSerde);

    MessageStream<KV<Integer, Integer>> inStream = appDesc.getInputStream(inputDescriptor);

    inStream.join(inStream, new TestJoinFunction(), integerSerde, kvSerde, kvSerde, JOIN_TTL, "join");
  }, config);

  createStreamOperatorTask(new SystemClock(), streamAppDesc); // should throw an exception
}
 

private StreamApplicationDescriptorImpl getAggregateTumblingWindowStreamGraph(AccumulationMode mode, Duration timeDuration,
      Trigger<IntegerEnvelope> earlyTrigger) throws IOException {
  StreamApplication userApp = appDesc -> {
    KVSerde<Integer, Integer> kvSerde = KVSerde.of(new IntegerSerde(), new IntegerSerde());
    GenericSystemDescriptor sd = new GenericSystemDescriptor("kafka", "mockFactoryClass");
    GenericInputDescriptor<KV<Integer, Integer>> inputDescriptor = sd.getInputDescriptor("integers", kvSerde);
    MessageStream<KV<Integer, Integer>> integers = appDesc.getInputStream(inputDescriptor);

    integers
        .map(new KVMapFunction())
        .window(Windows.<IntegerEnvelope, Integer>tumblingWindow(timeDuration, () -> 0, (m, c) -> c + 1, new IntegerSerde())
            .setEarlyTrigger(earlyTrigger)
            .setAccumulationMode(mode), "w1")
        .sink((message, messageCollector, taskCoordinator) -> {
          SystemStream outputSystemStream = new SystemStream("outputSystem", "outputStream");
          messageCollector.send(new OutgoingMessageEnvelope(outputSystemStream, message));
        });
  };

  return new StreamApplicationDescriptorImpl(userApp, config);
}
 

@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(expected = IllegalStateException.class)
public void testSetDefaultSystemDescriptorAfterGettingInputStream() {
  String streamId = "test-stream-1";
  GenericSystemDescriptor sd = new GenericSystemDescriptor("mockSystem", "mockSystemFactoryClass");
  GenericInputDescriptor isd = sd.getInputDescriptor(streamId, mock(Serde.class));

  new StreamApplicationDescriptorImpl(appDesc -> {
    appDesc.getInputStream(isd);
    appDesc.withDefaultSystem(sd); // should throw exception
  }, getConfig());
}
 

@Test
public void testGetInputStreamWithRelaxedTypes() {
  String streamId = "test-stream-1";
  GenericSystemDescriptor sd = new GenericSystemDescriptor("mockSystem", "mockSystemFactoryClass");
  GenericInputDescriptor isd = sd.getInputDescriptor(streamId, mock(Serde.class));
  StreamApplicationDescriptorImpl streamAppDesc = new StreamApplicationDescriptorImpl(appDesc -> {
    appDesc.getInputStream(isd);
  }, getConfig());

  InputOperatorSpec inputOpSpec = streamAppDesc.getInputOperators().get(streamId);
  assertEquals(OpCode.INPUT, inputOpSpec.getOpCode());
  assertEquals(streamId, inputOpSpec.getStreamId());
  assertEquals(isd, streamAppDesc.getInputDescriptors().get(streamId));
}
 

@Test(expected = IllegalArgumentException.class)
public void testGetInputStreamWithNullSerde() {
  GenericSystemDescriptor sd = new GenericSystemDescriptor("mockSystem", "mockSystemFactoryClass");
  GenericInputDescriptor isd = sd.getInputDescriptor("mockStreamId", null);
  new StreamApplicationDescriptorImpl(appDesc -> {
    appDesc.getInputStream(isd);
  }, getConfig());
}
 

@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 testOperatorGraphInitAndClose() {
  String inputStreamId1 = "input1";
  String inputStreamId2 = "input2";
  String inputSystem = "input-system";

  TaskName mockTaskName = mock(TaskName.class);
  TaskModel taskModel = mock(TaskModel.class);
  when(taskModel.getTaskName()).thenReturn(mockTaskName);
  when(this.context.getTaskContext().getTaskModel()).thenReturn(taskModel);

  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);

    Function mapFn = (Function & Serializable) m -> m;
    inputStream1.map(new TestMapFunction<Object, Object>("1", mapFn))
        .map(new TestMapFunction<Object, Object>("2", mapFn));

    inputStream2.map(new TestMapFunction<Object, Object>("3", mapFn))
        .map(new TestMapFunction<Object, Object>("4", mapFn));
  }, getConfig());

  OperatorImplGraph opImplGraph = new OperatorImplGraph(graphSpec.getOperatorSpecGraph(), this.context, SystemClock.instance());

  List<String> initializedOperators = BaseTestFunction.getInitListByTaskName(mockTaskName);

  // Assert that initialization occurs in topological order.
  assertEquals(initializedOperators.get(0), "1");
  assertEquals(initializedOperators.get(1), "2");
  assertEquals(initializedOperators.get(2), "3");
  assertEquals(initializedOperators.get(3), "4");

  // Assert that finalization occurs in reverse topological order.
  opImplGraph.close();
  List<String> closedOperators = BaseTestFunction.getCloseListByTaskName(mockTaskName);
  assertEquals(closedOperators.get(0), "4");
  assertEquals(closedOperators.get(1), "3");
  assertEquals(closedOperators.get(2), "2");
  assertEquals(closedOperators.get(3), "1");
}
 

@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));

  // verify that join function is initialized once.
  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);

  // verify that left partial join operator calls getFirstKey
  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));

  // verify that right partial join operator calls getSecondKey
  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));


  // verify that the join function apply is called with the correct messages on match
  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);
}
 

@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));
}
 

@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());
}