@VisibleForTesting
    public static JobGraph createJobGraph(StreamGraph streamGraph) {
        return new StreamingJobGraphGenerator(
                        Thread.currentThread().getContextClassLoader(),
                        streamGraph,
                        null,
                        Runnable::run)
                .createJobGraph();
    }
@Test
    void testPartitionTypesInBatchMode() {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setRuntimeMode(RuntimeExecutionMode.BATCH);
        env.setParallelism(4);
        env.disableOperatorChaining();
        DataStream<Integer> source = env.fromData(1);
        source
                // set the same parallelism as the source to make it a FORWARD exchange
                .map(value -> value)
                .setParallelism(1)
                .rescale()
                .map(value -> value)
                .rebalance()
                .map(value -> value)
                .keyBy(value -> value)
                .map(value -> value)
                .sinkTo(new DiscardingSink<>());

        JobGraph jobGraph = StreamingJobGraphGenerator.createJobGraph(env.getStreamGraph());
        List<JobVertex> verticesSorted = jobGraph.getVerticesSortedTopologicallyFromSources();
        assertHasOutputPartitionType(
                verticesSorted.get(0) /* source - forward */, ResultPartitionType.BLOCKING);
        assertHasOutputPartitionType(
                verticesSorted.get(1) /* rescale */, ResultPartitionType.BLOCKING);
        assertHasOutputPartitionType(
                verticesSorted.get(2) /* rebalance */, ResultPartitionType.BLOCKING);
        assertHasOutputPartitionType(
                verticesSorted.get(3) /* keyBy */, ResultPartitionType.BLOCKING);
        assertHasOutputPartitionType(
                verticesSorted.get(4) /* forward - sink */, ResultPartitionType.BLOCKING);
    }