samza-core/src/test/java/org/apache/samza/execution/TestJobGraphJsonGenerator.java - samza - Git at Google

 /*
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
  * regarding copyright ownership.  The ASF licenses this file
  * to you under the Apache License, Version 2.0 (the
  * "License"); you may not use this file except in compliance
  * with the License.  You may obtain a copy of the License at
  *
  *   http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing,
  * software distributed under the License is distributed on an
  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  * KIND, either express or implied.  See the License for the
  * specific language governing permissions and limitations
  * under the License.
  */

 package org.apache.samza.execution;

 import java.time.Duration;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
 import org.apache.samza.application.descriptors.StreamApplicationDescriptorImpl;
 import org.apache.samza.config.ApplicationConfig;
 import org.apache.samza.config.Config;
 import org.apache.samza.config.JobConfig;
 import org.apache.samza.config.MapConfig;
 import org.apache.samza.operators.functions.StreamTableJoinFunction;
 import org.apache.samza.operators.spec.OperatorSpec;
 import org.apache.samza.operators.spec.OperatorSpecs;
 import org.apache.samza.operators.spec.SendToTableOperatorSpec;
 import org.apache.samza.operators.spec.StreamTableJoinOperatorSpec;
 import org.apache.samza.system.descriptors.GenericInputDescriptor;
 import org.apache.samza.system.descriptors.GenericOutputDescriptor;
 import org.apache.samza.system.descriptors.GenericSystemDescriptor;
 import org.apache.samza.operators.KV;
 import org.apache.samza.operators.MessageStream;
 import org.apache.samza.operators.OutputStream;
 import org.apache.samza.operators.functions.JoinFunction;
 import org.apache.samza.operators.windows.Windows;
 import org.apache.samza.serializers.JsonSerdeV2;
 import org.apache.samza.serializers.KVSerde;
 import org.apache.samza.serializers.LongSerde;
 import org.apache.samza.serializers.NoOpSerde;
 import org.apache.samza.serializers.Serde;
 import org.apache.samza.serializers.StringSerde;
 import org.apache.samza.system.StreamSpec;
 import org.apache.samza.system.SystemAdmin;
 import org.apache.samza.system.SystemAdmins;
 import org.apache.samza.table.descriptors.TableDescriptor;
 import org.apache.samza.table.descriptors.TestLocalTableDescriptor;
 import org.apache.samza.testUtils.StreamTestUtils;
 import org.codehaus.jackson.map.ObjectMapper;
 import org.hamcrest.Matchers;
 import org.junit.Before;
 import org.junit.Test;

 import static org.apache.samza.execution.TestExecutionPlanner.*;
 import static org.junit.Assert.*;
 import static org.mockito.Mockito.*;


 /**
  * Unit test for {@link JobGraphJsonGenerator}
  */
 public class TestJobGraphJsonGenerator {
   private Config mockConfig;
   private JobNode mockJobNode;
   private StreamSpec input1Spec;
   private StreamSpec input2Spec;
   private StreamSpec outputSpec;
   private StreamSpec repartitionSpec;
   private KVSerde<String, Object> defaultSerde;
   private GenericSystemDescriptor inputSystemDescriptor;
   private GenericSystemDescriptor outputSystemDescriptor;
   private GenericSystemDescriptor intermediateSystemDescriptor;
   private GenericInputDescriptor<KV<String, Object>> input1Descriptor;
   private GenericInputDescriptor<KV<String, Object>> input2Descriptor;
   private GenericOutputDescriptor<KV<String, Object>> outputDescriptor;
   private TableDescriptor<String, Object, ?> table1Descriptor;
   private TableDescriptor<String, Object, ?> table2Descriptor;

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

   @Test
   public void testTaskApplication() throws Exception {
     JobGraphJsonGenerator jsonGenerator = new JobGraphJsonGenerator();
     JobGraph mockJobGraph = mock(JobGraph.class);
     ApplicationConfig mockAppConfig = mock(ApplicationConfig.class);
     when(mockAppConfig.getAppName()).thenReturn("testTaskApp");
     when(mockAppConfig.getAppId()).thenReturn("testTaskAppId");
     when(mockJobGraph.getApplicationConfig()).thenReturn(mockAppConfig);
     // compute the three disjoint sets of the JobGraph: input only, output only, and intermediate streams
     Set<StreamEdge> inEdges = new HashSet<>(mockJobNode.getInEdges().values());
     Set<StreamEdge> outEdges = new HashSet<>(mockJobNode.getOutEdges().values());
     Set<StreamEdge> intermediateEdges = new HashSet<>(inEdges);
     // intermediate streams are the intersection between input and output
     intermediateEdges.retainAll(outEdges);
     // remove all intermediate streams from input
     inEdges.removeAll(intermediateEdges);
     // remove all intermediate streams from output
     outEdges.removeAll(intermediateEdges);
     // set the return values for mockJobGraph
     when(mockJobGraph.getInputStreams()).thenReturn(inEdges);
     when(mockJobGraph.getOutputStreams()).thenReturn(outEdges);
     when(mockJobGraph.getIntermediateStreamEdges()).thenReturn(intermediateEdges);
     Set<TableDescriptor> tables = new HashSet<>(mockJobNode.getTables().values());
     when(mockJobGraph.getTables()).thenReturn(tables);
     when(mockJobGraph.getJobNodes()).thenReturn(Collections.singletonList(mockJobNode));
     String graphJson = jsonGenerator.toJson(mockJobGraph);
     ObjectMapper objectMapper = new ObjectMapper();
     JobGraphJsonGenerator.JobGraphJson jsonObject = objectMapper.readValue(graphJson.getBytes(), JobGraphJsonGenerator.JobGraphJson.class);
     assertEquals("testTaskAppId", jsonObject.applicationId);
     assertEquals("testTaskApp", jsonObject.applicationName);
     Set<String> inStreamIds = inEdges.stream().map(stream -> stream.getStreamSpec().getId()).collect(Collectors.toSet());
     assertThat(jsonObject.sourceStreams.keySet(), Matchers.containsInAnyOrder(inStreamIds.toArray()));
     Set<String> outStreamIds = outEdges.stream().map(stream -> stream.getStreamSpec().getId()).collect(Collectors.toSet());
     assertThat(jsonObject.sinkStreams.keySet(), Matchers.containsInAnyOrder(outStreamIds.toArray()));
     Set<String> intStreamIds = intermediateEdges.stream().map(stream -> stream.getStreamSpec().getId()).collect(Collectors.toSet());
     assertThat(jsonObject.intermediateStreams.keySet(), Matchers.containsInAnyOrder(intStreamIds.toArray()));
     Set<String> tableIds = tables.stream().map(t -> t.getTableId()).collect(Collectors.toSet());
     assertThat(jsonObject.tables.keySet(), Matchers.containsInAnyOrder(tableIds.toArray()));
     JobGraphJsonGenerator.JobNodeJson expectedNodeJson = new JobGraphJsonGenerator.JobNodeJson();
     expectedNodeJson.jobId = mockJobNode.getJobId();
     expectedNodeJson.jobName = mockJobNode.getJobName();
     assertEquals(1, jsonObject.jobs.size());
     JobGraphJsonGenerator.JobNodeJson actualNodeJson = jsonObject.jobs.get(0);
     assertEquals(expectedNodeJson.jobId, actualNodeJson.jobId);
     assertEquals(expectedNodeJson.jobName, actualNodeJson.jobName);
     assertEquals(3, actualNodeJson.operatorGraph.inputStreams.size());
     assertEquals(2, actualNodeJson.operatorGraph.outputStreams.size());
     assertEquals(0, actualNodeJson.operatorGraph.operators.size());
   }

   @Test
   public void testLegacyTaskApplication() throws Exception {
     JobGraphJsonGenerator jsonGenerator = new JobGraphJsonGenerator();
     JobGraph mockJobGraph = mock(JobGraph.class);
     ApplicationConfig mockAppConfig = mock(ApplicationConfig.class);
     when(mockAppConfig.getAppName()).thenReturn("testTaskApp");
     when(mockAppConfig.getAppId()).thenReturn("testTaskAppId");
     when(mockJobGraph.getApplicationConfig()).thenReturn(mockAppConfig);
     String graphJson = jsonGenerator.toJson(mockJobGraph);
     ObjectMapper objectMapper = new ObjectMapper();
     JobGraphJsonGenerator.JobGraphJson jsonObject = objectMapper.readValue(graphJson.getBytes(), JobGraphJsonGenerator.JobGraphJson.class);
     assertEquals("testTaskAppId", jsonObject.applicationId);
     assertEquals("testTaskApp", jsonObject.applicationName);
     JobGraphJsonGenerator.JobNodeJson expectedNodeJson = new JobGraphJsonGenerator.JobNodeJson();
     expectedNodeJson.jobId = mockJobNode.getJobId();
     expectedNodeJson.jobName = mockJobNode.getJobName();
     assertEquals(0, jsonObject.jobs.size());
   }

   public class PageViewEvent {
     String getCountry() {
       return "";
     }
   }

   @Test
   public void testOperatorToMapForTable() {
     JobGraphJsonGenerator jsonGenerator = new JobGraphJsonGenerator();
     Map<String, Object> map;
     SendToTableOperatorSpec<Object, Object> sendToTableOperatorSpec =
         OperatorSpecs.createSendToTableOperatorSpec("test-sent-to-table", "test-sent-to");
     map = jsonGenerator.operatorToMap(sendToTableOperatorSpec);
     assertTrue(map.containsKey("tableId"));
     assertEquals(map.get("tableId"), "test-sent-to-table");
     assertEquals(map.get("opCode"), OperatorSpec.OpCode.SEND_TO.name());
     assertEquals(map.get("opId"), "test-sent-to");
     StreamTableJoinOperatorSpec<String, String, String, String> streamTableJoinOperatorSpec =
         OperatorSpecs.createStreamTableJoinOperatorSpec("test-join-table", mock(StreamTableJoinFunction.class), "test-join");
     map = jsonGenerator.operatorToMap(streamTableJoinOperatorSpec);
     assertTrue(map.containsKey("tableId"));
     assertEquals(map.get("tableId"), "test-join-table");
     assertEquals(map.get("opCode"), OperatorSpec.OpCode.JOIN.name());
     assertEquals(map.get("opId"), "test-join");
   }

 }
	/*
	* Licensed to the Apache Software Foundation (ASF) under one
	* or more contributor license agreements. See the NOTICE file
	* distributed with this work for additional information
	* regarding copyright ownership. The ASF licenses this file
	* to you under the Apache License, Version 2.0 (the
	* "License"); you may not use this file except in compliance
	* with the License. You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing,
	* software distributed under the License is distributed on an
	* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
	* KIND, either express or implied. See the License for the
	* specific language governing permissions and limitations
	* under the License.
	*/

	package org.apache.samza.execution;

	import java.time.Duration;
	import java.util.Collections;
	import java.util.HashMap;
	import java.util.HashSet;
	import java.util.Map;
	import java.util.Set;
	import java.util.stream.Collectors;
	import org.apache.samza.application.descriptors.StreamApplicationDescriptorImpl;
	import org.apache.samza.config.ApplicationConfig;
	import org.apache.samza.config.Config;
	import org.apache.samza.config.JobConfig;
	import org.apache.samza.config.MapConfig;
	import org.apache.samza.operators.functions.StreamTableJoinFunction;
	import org.apache.samza.operators.spec.OperatorSpec;
	import org.apache.samza.operators.spec.OperatorSpecs;
	import org.apache.samza.operators.spec.SendToTableOperatorSpec;
	import org.apache.samza.operators.spec.StreamTableJoinOperatorSpec;
	import org.apache.samza.system.descriptors.GenericInputDescriptor;
	import org.apache.samza.system.descriptors.GenericOutputDescriptor;
	import org.apache.samza.system.descriptors.GenericSystemDescriptor;
	import org.apache.samza.operators.KV;
	import org.apache.samza.operators.MessageStream;
	import org.apache.samza.operators.OutputStream;
	import org.apache.samza.operators.functions.JoinFunction;
	import org.apache.samza.operators.windows.Windows;
	import org.apache.samza.serializers.JsonSerdeV2;
	import org.apache.samza.serializers.KVSerde;
	import org.apache.samza.serializers.LongSerde;
	import org.apache.samza.serializers.NoOpSerde;
	import org.apache.samza.serializers.Serde;
	import org.apache.samza.serializers.StringSerde;
	import org.apache.samza.system.StreamSpec;
	import org.apache.samza.system.SystemAdmin;
	import org.apache.samza.system.SystemAdmins;
	import org.apache.samza.table.descriptors.TableDescriptor;
	import org.apache.samza.table.descriptors.TestLocalTableDescriptor;
	import org.apache.samza.testUtils.StreamTestUtils;
	import org.codehaus.jackson.map.ObjectMapper;
	import org.hamcrest.Matchers;
	import org.junit.Before;
	import org.junit.Test;

	import static org.apache.samza.execution.TestExecutionPlanner.*;
	import static org.junit.Assert.*;
	import static org.mockito.Mockito.*;


	/**
	* Unit test for {@link JobGraphJsonGenerator}
	*/
	public class TestJobGraphJsonGenerator {
	private Config mockConfig;
	private JobNode mockJobNode;
	private StreamSpec input1Spec;
	private StreamSpec input2Spec;
	private StreamSpec outputSpec;
	private StreamSpec repartitionSpec;
	private KVSerde<String, Object> defaultSerde;
	private GenericSystemDescriptor inputSystemDescriptor;
	private GenericSystemDescriptor outputSystemDescriptor;
	private GenericSystemDescriptor intermediateSystemDescriptor;
	private GenericInputDescriptor<KV<String, Object>> input1Descriptor;
	private GenericInputDescriptor<KV<String, Object>> input2Descriptor;
	private GenericOutputDescriptor<KV<String, Object>> outputDescriptor;
	private TableDescriptor<String, Object, ?> table1Descriptor;
	private TableDescriptor<String, Object, ?> table2Descriptor;

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

	@Test
	public void testTaskApplication() throws Exception {
	JobGraphJsonGenerator jsonGenerator = new JobGraphJsonGenerator();
	JobGraph mockJobGraph = mock(JobGraph.class);
	ApplicationConfig mockAppConfig = mock(ApplicationConfig.class);
	when(mockAppConfig.getAppName()).thenReturn("testTaskApp");
	when(mockAppConfig.getAppId()).thenReturn("testTaskAppId");
	when(mockJobGraph.getApplicationConfig()).thenReturn(mockAppConfig);
	// compute the three disjoint sets of the JobGraph: input only, output only, and intermediate streams
	Set<StreamEdge> inEdges = new HashSet<>(mockJobNode.getInEdges().values());
	Set<StreamEdge> outEdges = new HashSet<>(mockJobNode.getOutEdges().values());
	Set<StreamEdge> intermediateEdges = new HashSet<>(inEdges);
	// intermediate streams are the intersection between input and output
	intermediateEdges.retainAll(outEdges);
	// remove all intermediate streams from input
	inEdges.removeAll(intermediateEdges);
	// remove all intermediate streams from output
	outEdges.removeAll(intermediateEdges);
	// set the return values for mockJobGraph
	when(mockJobGraph.getInputStreams()).thenReturn(inEdges);
	when(mockJobGraph.getOutputStreams()).thenReturn(outEdges);
	when(mockJobGraph.getIntermediateStreamEdges()).thenReturn(intermediateEdges);
	Set<TableDescriptor> tables = new HashSet<>(mockJobNode.getTables().values());
	when(mockJobGraph.getTables()).thenReturn(tables);
	when(mockJobGraph.getJobNodes()).thenReturn(Collections.singletonList(mockJobNode));
	String graphJson = jsonGenerator.toJson(mockJobGraph);
	ObjectMapper objectMapper = new ObjectMapper();
	JobGraphJsonGenerator.JobGraphJson jsonObject = objectMapper.readValue(graphJson.getBytes(), JobGraphJsonGenerator.JobGraphJson.class);
	assertEquals("testTaskAppId", jsonObject.applicationId);
	assertEquals("testTaskApp", jsonObject.applicationName);
	Set<String> inStreamIds = inEdges.stream().map(stream -> stream.getStreamSpec().getId()).collect(Collectors.toSet());
	assertThat(jsonObject.sourceStreams.keySet(), Matchers.containsInAnyOrder(inStreamIds.toArray()));
	Set<String> outStreamIds = outEdges.stream().map(stream -> stream.getStreamSpec().getId()).collect(Collectors.toSet());
	assertThat(jsonObject.sinkStreams.keySet(), Matchers.containsInAnyOrder(outStreamIds.toArray()));
	Set<String> intStreamIds = intermediateEdges.stream().map(stream -> stream.getStreamSpec().getId()).collect(Collectors.toSet());
	assertThat(jsonObject.intermediateStreams.keySet(), Matchers.containsInAnyOrder(intStreamIds.toArray()));
	Set<String> tableIds = tables.stream().map(t -> t.getTableId()).collect(Collectors.toSet());
	assertThat(jsonObject.tables.keySet(), Matchers.containsInAnyOrder(tableIds.toArray()));
	JobGraphJsonGenerator.JobNodeJson expectedNodeJson = new JobGraphJsonGenerator.JobNodeJson();
	expectedNodeJson.jobId = mockJobNode.getJobId();
	expectedNodeJson.jobName = mockJobNode.getJobName();
	assertEquals(1, jsonObject.jobs.size());
	JobGraphJsonGenerator.JobNodeJson actualNodeJson = jsonObject.jobs.get(0);
	assertEquals(expectedNodeJson.jobId, actualNodeJson.jobId);
	assertEquals(expectedNodeJson.jobName, actualNodeJson.jobName);
	assertEquals(3, actualNodeJson.operatorGraph.inputStreams.size());
	assertEquals(2, actualNodeJson.operatorGraph.outputStreams.size());
	assertEquals(0, actualNodeJson.operatorGraph.operators.size());
	}

	@Test
	public void testLegacyTaskApplication() throws Exception {
	JobGraphJsonGenerator jsonGenerator = new JobGraphJsonGenerator();
	JobGraph mockJobGraph = mock(JobGraph.class);
	ApplicationConfig mockAppConfig = mock(ApplicationConfig.class);
	when(mockAppConfig.getAppName()).thenReturn("testTaskApp");
	when(mockAppConfig.getAppId()).thenReturn("testTaskAppId");
	when(mockJobGraph.getApplicationConfig()).thenReturn(mockAppConfig);
	String graphJson = jsonGenerator.toJson(mockJobGraph);
	ObjectMapper objectMapper = new ObjectMapper();
	JobGraphJsonGenerator.JobGraphJson jsonObject = objectMapper.readValue(graphJson.getBytes(), JobGraphJsonGenerator.JobGraphJson.class);
	assertEquals("testTaskAppId", jsonObject.applicationId);
	assertEquals("testTaskApp", jsonObject.applicationName);
	JobGraphJsonGenerator.JobNodeJson expectedNodeJson = new JobGraphJsonGenerator.JobNodeJson();
	expectedNodeJson.jobId = mockJobNode.getJobId();
	expectedNodeJson.jobName = mockJobNode.getJobName();
	assertEquals(0, jsonObject.jobs.size());
	}

	public class PageViewEvent {
	String getCountry() {
	return "";
	}
	}

	@Test
	public void testOperatorToMapForTable() {
	JobGraphJsonGenerator jsonGenerator = new JobGraphJsonGenerator();
	Map<String, Object> map;
	SendToTableOperatorSpec<Object, Object> sendToTableOperatorSpec =
	OperatorSpecs.createSendToTableOperatorSpec("test-sent-to-table", "test-sent-to");
	map = jsonGenerator.operatorToMap(sendToTableOperatorSpec);
	assertTrue(map.containsKey("tableId"));
	assertEquals(map.get("tableId"), "test-sent-to-table");
	assertEquals(map.get("opCode"), OperatorSpec.OpCode.SEND_TO.name());
	assertEquals(map.get("opId"), "test-sent-to");
	StreamTableJoinOperatorSpec<String, String, String, String> streamTableJoinOperatorSpec =
	OperatorSpecs.createStreamTableJoinOperatorSpec("test-join-table", mock(StreamTableJoinFunction.class), "test-join");
	map = jsonGenerator.operatorToMap(streamTableJoinOperatorSpec);
	assertTrue(map.containsKey("tableId"));
	assertEquals(map.get("tableId"), "test-join-table");
	assertEquals(map.get("opCode"), OperatorSpec.OpCode.JOIN.name());
	assertEquals(map.get("opId"), "test-join");
	}

	}