@Override
    public void open() throws Exception {
        this.timerService =
                getInternalTimerService("processing timer", VoidNamespaceSerializer.INSTANCE, this);
        this.keySet = new HashSet<>();
        super.open();
    }
@Test
    void testEndInput() throws Exception {
        AtomicInteger counter = new AtomicInteger();
        KeyedProcessOperator<Integer, Integer, Integer> processOperator =
                new KeyedProcessOperator<>(
                        new OneInputStreamProcessFunction<Integer, Integer>() {
                            @Override
                            public void processRecord(
                                    Integer record,
                                    Collector<Integer> output,
                                    PartitionedContext ctx) {
                                // do nothing.
                            }

                            @Override
                            public void endInput(NonPartitionedContext<Integer> ctx) {
                                try {
                                    ctx.applyToAllPartitions(
                                            (out, context) -> {
                                                counter.incrementAndGet();
                                                Integer currentKey =
                                                        context.getStateManager().getCurrentKey();
                                                out.collect(currentKey);
                                            });
                                } catch (Exception e) {
                                    throw new RuntimeException(e);
                                }
                            }
                        });

        try (KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness =
                new KeyedOneInputStreamOperatorTestHarness<>(
                        processOperator,
                        (KeySelector<Integer, Integer>) value -> value,
                        Types.INT)) {
            testHarness.open();
            testHarness.processElement(new StreamRecord<>(1)); // key is 1
            testHarness.processElement(new StreamRecord<>(2)); // key is 2
            testHarness.processElement(new StreamRecord<>(3)); // key is 3
            testHarness.endInput();
            assertThat(counter).hasValue(3);
            Collection<StreamRecord<Integer>> recordOutput = testHarness.getRecordOutput();
            assertThat(recordOutput)
                    .containsExactly(
                            new StreamRecord<>(1), new StreamRecord<>(2), new StreamRecord<>(3));
        }
    }