Repository: apache/flink
File: flink-runtime/src/test/java/org/apache/flink/runtime/jobmaster/JobMasterTest.java (Java)

Code before:
                testingResourceManagerGateway.getAddress(), testingResourceManagerGateway);
        return testingResourceManagerGateway;
    }

    /**
     * Tests that the job execution is failed if the TaskExecutor disconnects from the JobMaster.
     */
    @Test
    void testJobFailureWhenGracefulTaskExecutorTermination() throws Exception {
        final List<Event> jobEvents = new ArrayList<>();
        runJobFailureWhenTaskExecutorTerminatesTest(
                heartbeatServices,
                (localTaskManagerLocation, jobMasterGateway) ->
                        jobMasterGateway.disconnectTaskManager(
                                localTaskManagerLocation.getResourceID(),
                                new FlinkException("Test disconnectTaskManager exception.")),
                jobEvents);
        assertThat(
                        jobEvents.stream()
                                .filter(
                                        event ->
                                                Events.JobStatusChangeEvent.name()
                                                        .equals(event.getName()))
                                .map(Event::getAttributes)
                                .map(x -> x.get("newJobStatus")))
                .containsSubsequence(
                        JobStatus.RUNNING.toString(),
                        JobStatus.FAILING.toString(),
                        JobStatus.FAILED.toString());

        assertThat(
                        jobEvents.stream()
                                .filter(
                                        event ->
                                                Events.AllSubtasksStatusChangeEvent.name()
                                                        .equals(event.getName()))
                                .map(Event::getAttributes)
                                .map(
                                        x ->
                                                x.get(
                                                        AllSubTasksRunningOrFinishedStateTimeMetrics
                                                                .STATUS_ATTRIBUTE)))
                .containsExactly(
                        ALL_RUNNING_OR_FINISHED.toString(), NOT_ALL_RUNNING_OR_FINISHED.toString());
    }

    @Test
    void testJobFailureWhenTaskExecutorHeartbeatTimeout() throws Exception {
        final List<Event> jobEvents = new ArrayList<>();
        final TestingHeartbeatServices testingHeartbeatService =
                new TestingHeartbeatServices(heartbeatInterval, heartbeatTimeout);

Review comment (bug):
I had a quick look at this is. My concern is that the CREATED event might not be coming out - and this is a Flink bug. I would feel for confident if the test checks that the CREATED event is coming out and the code be arranged so that the race condition does not occur. WDYT?

Code after:
    @Test
    void testJobFailureWhenGracefulTaskExecutorTermination() throws Exception {
        final List<Event> jobEvents = new ArrayList<>();
        final SchedulerType schedulerType =
                runJobFailureWhenTaskExecutorTerminatesTest(
                        heartbeatServices,
                        (localTaskManagerLocation, jobMasterGateway) ->
                                jobMasterGateway.disconnectTaskManager(
                                        localTaskManagerLocation.getResourceID(),
                                        new FlinkException(
                                                "Test disconnectTaskManager exception.")),
                        jobEvents);

        assertJobStatusTransitions(schedulerType, jobEvents);
    }

    @Test
    void testJobFailureWhenTaskExecutorHeartbeatTimeout() throws Exception {
        final List<Event> jobEvents = new ArrayList<>();
        final TestingHeartbeatServices testingHeartbeatService =
                new TestingHeartbeatServices(heartbeatInterval, heartbeatTimeout);

        final SchedulerType schedulerType =
                runJobFailureWhenTaskExecutorTerminatesTest(
                        testingHeartbeatService,
                        (localTaskManagerLocation, jobMasterGateway) ->
                                testingHeartbeatService.triggerHeartbeatTimeout(
                                        jmResourceId, localTaskManagerLocation.getResourceID()),
                        jobEvents);

        assertJobStatusTransitions(schedulerType, jobEvents);
    }

    /**
     * Asserts that job status transitions are as expected based on the scheduler type.
     * DefaultScheduler does not emit CREATED state, while AdaptiveScheduler and
     * AdaptiveBatchScheduler do.
     */
    private static void assertJobStatusTransitions(
            SchedulerType schedulerType, List<Event> jobEvents) {
        final List<String> jobStatusTransitions =
                jobEvents.stream()
                        .filter(event -> Events.JobStatusChangeEvent.name().equals(event.getName()))
                        .map(Event::getAttributes)
                        .map(x -> (String) x.get("newJobStatus"))
                        .collect(Collectors.toList());

        if (schedulerType == SchedulerType.Adaptive) {
            // Adaptive schedulers emit CREATED: CREATED → RUNNING → FAILING → FAILED
            assertThat(jobStatusTransitions)
                    .containsExactly(