use of org.apache.flink.api.java.io.TypeSerializerInputFormat in project flink by apache.
the class StreamingJobGraphGeneratorTest method testInputOutputFormat.
@Test
public void testInputOutputFormat() {
StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
DataStream<Long> source = env.addSource(new InputFormatSourceFunction<>(new TypeSerializerInputFormat<>(TypeInformation.of(Long.class)), TypeInformation.of(Long.class)), TypeInformation.of(Long.class)).name("source");
source.writeUsingOutputFormat(new DiscardingOutputFormat<>()).name("sink1");
source.writeUsingOutputFormat(new DiscardingOutputFormat<>()).name("sink2");
StreamGraph streamGraph = env.getStreamGraph();
JobGraph jobGraph = StreamingJobGraphGenerator.createJobGraph(streamGraph);
assertEquals(1, jobGraph.getNumberOfVertices());
JobVertex jobVertex = jobGraph.getVertices().iterator().next();
assertTrue(jobVertex instanceof InputOutputFormatVertex);
InputOutputFormatContainer formatContainer = new InputOutputFormatContainer(new TaskConfig(jobVertex.getConfiguration()), Thread.currentThread().getContextClassLoader());
Map<OperatorID, UserCodeWrapper<? extends InputFormat<?, ?>>> inputFormats = formatContainer.getInputFormats();
Map<OperatorID, UserCodeWrapper<? extends OutputFormat<?>>> outputFormats = formatContainer.getOutputFormats();
assertEquals(1, inputFormats.size());
assertEquals(2, outputFormats.size());
Map<String, OperatorID> nameToOperatorIds = new HashMap<>();
StreamConfig headConfig = new StreamConfig(jobVertex.getConfiguration());
nameToOperatorIds.put(headConfig.getOperatorName(), headConfig.getOperatorID());
Map<Integer, StreamConfig> chainedConfigs = headConfig.getTransitiveChainedTaskConfigs(Thread.currentThread().getContextClassLoader());
for (StreamConfig config : chainedConfigs.values()) {
nameToOperatorIds.put(config.getOperatorName(), config.getOperatorID());
}
InputFormat<?, ?> sourceFormat = inputFormats.get(nameToOperatorIds.get("Source: source")).getUserCodeObject();
assertTrue(sourceFormat instanceof TypeSerializerInputFormat);
OutputFormat<?> sinkFormat1 = outputFormats.get(nameToOperatorIds.get("Sink: sink1")).getUserCodeObject();
assertTrue(sinkFormat1 instanceof DiscardingOutputFormat);
OutputFormat<?> sinkFormat2 = outputFormats.get(nameToOperatorIds.get("Sink: sink2")).getUserCodeObject();
assertTrue(sinkFormat2 instanceof DiscardingOutputFormat);
}
Aggregations