Repository: apache/flink
File: flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/async/AsyncWaitOperatorTest.java (Java)

Code before:
        StreamTaskMailboxTestHarnessBuilder<Integer> builder =
                new StreamTaskMailboxTestHarnessBuilder<>(
                                OneInputStreamTask::new, BasicTypeInfo.INT_TYPE_INFO)
                        .addInput(BasicTypeInfo.INT_TYPE_INFO);

        try (StreamTaskMailboxTestHarness<Integer> testHarness =
                builder.setupOutputForSingletonOperatorChain(
                                new AsyncWaitOperatorFactory<Integer, Integer>(
                                        new AlwaysTimeoutAsyncFunction(),
                                        TIMEOUT,
                                        10,
                                        mode,
                                        exceptionRetryStrategy))
                        .build()) {

            final long initialTime = 0L;
            AtomicReference<Throwable> error = new AtomicReference<>();
            testHarness.getStreamMockEnvironment().setExternalExceptionHandler(error::set);

            try {
                testHarness.processElement(new StreamRecord<>(1, initialTime + 1));
                testHarness.processElement(new StreamRecord<>(2, initialTime + 2));
                while (error.get() == null) {
                    testHarness.processAll();
                }
            } catch (Exception e) {
                error.set(e);
            }
            ExceptionUtils.assertThrowableWithMessage(error.get(), "Dummy timeout error");
        }
    }

    @Test
    public void testProcessingTimeWithMailboxThreadOrdered() throws Exception {
        testProcessingTimeWithCollectFromMailboxThread(
                AsyncDataStream.OutputMode.ORDERED, NO_RETRY_STRATEGY);
    }

    @Test
    public void testProcessingTimeWithMailboxThreadUnordered() throws Exception {
        testProcessingTimeWithCollectFromMailboxThread(
                AsyncDataStream.OutputMode.UNORDERED, NO_RETRY_STRATEGY);
    }

    @Test
    public void testProcessingTimeWithMailboxThreadOrderedWithRetry() throws Exception {
        testProcessingTimeWithCollectFromMailboxThread(
                AsyncDataStream.OutputMode.ORDERED, exceptionRetryStrategy);
    }

    @Test

Review comment (suggestion):
If you counted the number of invocations, you could verify that it only ever tries once, even after the timeout exception.

Code after:
        StreamTaskMailboxTestHarnessBuilder<Integer> builder =
                new StreamTaskMailboxTestHarnessBuilder<>(
                                OneInputStreamTask::new, BasicTypeInfo.INT_TYPE_INFO)
                        .addInput(BasicTypeInfo.INT_TYPE_INFO);

        AlwaysTimeoutAsyncFunction asyncFunction = new AlwaysTimeoutAsyncFunction();
        try (StreamTaskMailboxTestHarness<Integer> testHarness =
                builder.setupOutputForSingletonOperatorChain(
                                new AsyncWaitOperatorFactory<Integer, Integer>(
                                        asyncFunction, TIMEOUT, 10, mode, exceptionRetryStrategy))
                        .build()) {

            final long initialTime = 0L;
            AtomicReference<Throwable> error = new AtomicReference<>();
            testHarness.getStreamMockEnvironment().setExternalExceptionHandler(error::set);

            try {
                testHarness.processElement(new StreamRecord<>(1, initialTime + 1));
                testHarness.processElement(new StreamRecord<>(2, initialTime + 2));
                while (error.get() == null) {
                    testHarness.processAll();
                }
            } catch (Exception e) {
                error.set(e);
            }
            ExceptionUtils.assertThrowableWithMessage(error.get(), "Dummy timeout error");
            // verify the 1st element's try count is exactly 1
            assertThat(asyncFunction.getTryCount(1)).isEqualTo(1);
        }
    }

    @Test
    public void testProcessingTimeWithMailboxThreadOrdered() throws Exception {
        testProcessingTimeWithCollectFromMailboxThread(
                AsyncDataStream.OutputMode.ORDERED, NO_RETRY_STRATEGY);
    }

    @Test
    public void testProcessingTimeWithMailboxThreadUnordered() throws Exception {
        testProcessingTimeWithCollectFromMailboxThread(
                AsyncDataStream.OutputMode.UNORDERED, NO_RETRY_STRATEGY);
    }

    @Test
    public void testProcessingTimeWithMailboxThreadOrderedWithRetry() throws Exception {
        testProcessingTimeWithCollectFromMailboxThread(
                AsyncDataStream.OutputMode.ORDERED, exceptionRetryStrategy);
    }

    @Test
    public void testProcessingTimeWithMailboxThreadUnorderedWithRetry() throws Exception {