Repository: apache/flink
File: flink-table/flink-table-runtime/src/test/java/org/apache/flink/table/runtime/operators/join/lookup/TableKeyedAsyncWaitOperatorTest.java (Java)

Code before:
                expected,
                testHarness.getOutput(),
                new StreamRecordComparator());
        testHarness.close();

        testLazyAsyncFunction = new LazyAsyncFunction();
        testLazyAsyncFunction.countDown();
        testHarness = createKeyedTestHarness(testLazyAsyncFunction, TIMEOUT, 10);

        testHarness.initializeState(snapshot);
        testHarness.open();

        testHarness.processElement(new StreamRecord<>(5, initialTime + 5));
        testHarness.processElement(new StreamRecord<>(6, initialTime + 6));
        testHarness.processElement(new StreamRecord<>(7, initialTime + 7));
        testHarness.processElement(new StreamRecord<>(8, initialTime + 8));
        testHarness.processWatermark(Long.MAX_VALUE);

        expected.add(new StreamRecord<>(5, initialTime + 5));
        expected.add(new StreamRecord<>(6, initialTime + 6));
        expected.add(new StreamRecord<>(7, initialTime + 7));
        expected.add(new StreamRecord<>(8, initialTime + 8));
        expected.add(new Watermark(Long.MAX_VALUE));

        testLazyAsyncFunction.countDown();
        testHarness.endInput();

        TestHarnessUtil.assertOutputEqualsSorted(
                "StateAndRestored Test Output was not correct.",
                expected,
                testHarness.getOutput(),
                new StreamRecordComparator());
        testHarness.close();
    }

    @Test
    public void testKeyedAsyncTimeoutFailure() throws Exception {
        testKeyedAsyncTimeout(
                new LazyAsyncFunction(),
                Optional.of(TimeoutException.class),
                List.of(new StreamRecord<>(2, 5L)));
    }

    @Test
    public void testKeyedAsyncTimeoutIgnore() throws Exception {
        testKeyedAsyncTimeout(
                new IgnoreTimeoutLazyAsyncFunction(),
                Optional.empty(),
                List.of(new StreamRecord<>(3, 0L), new StreamRecord<>(2, 5L)));
    }


Review comment (suggestion):
Maybe we can move `processWatermark` between `initializeState()` and `StreamRecord(5, 5)`, so we can check if ongoingRecordCount meets out expectations. 

In this case, the ongoingRecordCount of activeEpoch should change 4 -> 0 -> 4 -> 0 when executing `before processWatermark` -> `after processWatermark` -> `add StreamRecord(5, 6, 7, 8)` -> `endInput`

Code after:
                "StateAndRestored Test Output was not correct before restore.",
                expected,
                testHarness.getOutput(),
                new StreamRecordComparator());
        testHarness.close();

        testLazyAsyncFunction = new LazyAsyncFunction();
        testHarness = createKeyedTestHarness(testLazyAsyncFunction, TIMEOUT, 10);

        testHarness.initializeState(snapshot);
        testHarness.open();

        Optional<Epoch<Integer>> epochWithMinWatermark = unwrapEpoch(testHarness, Long.MIN_VALUE);
        assertThat(epochWithMinWatermark).isPresent();
        assertThat(epochWithMinWatermark.get().getOngoingRecordCount()).isEqualTo(4);

        testHarness.processWatermark(1000L);
        testHarness.processElement(new StreamRecord<>(5, initialTime + 5));
        testHarness.processElement(new StreamRecord<>(6, initialTime + 6));
        testHarness.processElement(new StreamRecord<>(7, initialTime + 7));
        testHarness.processElement(new StreamRecord<>(8, initialTime + 8));
        Optional<Epoch<Integer>> epochWithMidWatermark = unwrapEpoch(testHarness, 1000L);
        assertThat(epochWithMidWatermark).isPresent();
        assertThat(epochWithMidWatermark.get().getOngoingRecordCount()).isEqualTo(4);

        testHarness.processWatermark(Long.MAX_VALUE);
        Optional<Epoch<Integer>> epochWithMaxWatermark = unwrapEpoch(testHarness, Long.MAX_VALUE);
        assertThat(epochWithMaxWatermark).isPresent();
        assertThat(epochWithMaxWatermark.get().getOngoingRecordCount()).isEqualTo(0);

        expected.add(new Watermark(1000L));
        expected.add(new StreamRecord<>(5, initialTime + 5));
        expected.add(new StreamRecord<>(6, initialTime + 6));
        expected.add(new StreamRecord<>(7, initialTime + 7));
        expected.add(new StreamRecord<>(8, initialTime + 8));
        expected.add(new Watermark(Long.MAX_VALUE));

        testLazyAsyncFunction.countDown();
        testHarness.endInput();

        // TODO FLINK-37981 send Long.MAX watermark downstream
        assertThat(unwrapEpochManager(testHarness).getOutputQueue().size()).isEqualTo(1);
        assertThat(
                        unwrapEpochManager(testHarness)
                                .getOutputQueue()
                                .get(0)
                                .getWatermark()
                                .getTimestamp())
                .isEqualTo(Long.MAX_VALUE);

        TestHarnessUtil.assertOutputEqualsSorted(