@Override
    public void endInput(int inputId) throws Exception {
        // sanity check.
        checkState(inputId >= 1 && inputId <= 2);
        if (inputId == 1) {
            userFunction.endNonBroadcastInput(nonPartitionedContext);
        } else {
            userFunction.endBroadcastInput(nonPartitionedContext);
        }
    }
@Test
    void testEndInput() throws Exception {
        AtomicInteger nonBroadcastInputCounter = new AtomicInteger();
        AtomicInteger broadcastInputCounter = new AtomicInteger();
        TwoInputBroadcastProcessOperator<Integer, Long, Long> processOperator =
                new TwoInputBroadcastProcessOperator<>(
                        new TwoInputBroadcastStreamProcessFunction<Integer, Long, Long>() {

                            @Override
                            public void processRecordFromNonBroadcastInput(
                                    Integer record,
                                    Collector<Long> output,
                                    PartitionedContext ctx) {
                                // do nothing.
                            }

                            @Override
                            public void processRecordFromBroadcastInput(
                                    Long record, NonPartitionedContext<Long> ctx) {
                                // do nothing.
                            }

                            @Override
                            public void endNonBroadcastInput(NonPartitionedContext<Long> ctx) {
                                try {
                                    ctx.applyToAllPartitions(
                                            (out, context) -> {
                                                nonBroadcastInputCounter.incrementAndGet();
                                                out.collect(1L);
                                            });
                                } catch (Exception e) {
                                    throw new RuntimeException(e);
                                }
                            }

                            @Override
                            public void endBroadcastInput(NonPartitionedContext<Long> ctx) {
                                try {
                                    ctx.applyToAllPartitions(
                                            (out, context) -> {
                                                broadcastInputCounter.incrementAndGet();
                                                out.collect(2L);
                                            });
                                } catch (Exception e) {
                                    throw new RuntimeException(e);
                                }
                            }
                        });

        try (TwoInputStreamOperatorTestHarness<Integer, Long, Long> testHarness =
                new TwoInputStreamOperatorTestHarness<>(processOperator)) {
            testHarness.open();
            testHarness.endInput1();
            assertThat(nonBroadcastInputCounter).hasValue(1);
            testHarness.endInput2();
            assertThat(broadcastInputCounter).hasValue(1);
            Collection<StreamRecord<Long>> recordOutput = testHarness.getRecordOutput();
            assertThat(recordOutput)
                    .containsExactly(new StreamRecord<>(1L), new StreamRecord<>(2L));
        }
    }