From 46c7d366432e2e054e53ff526bc74169f022d5c1 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Wed, 1 Apr 2026 17:51:01 -0700 Subject: [PATCH] fix flaky DAGSchedulerSuite SPARK-40082 test for pushBasedShuffle Add sc.listenerBus.waitUntilEmpty() before the completedStage assertion to ensure the async listener bus has processed the stage completion event before we check it. Co-authored-by: DB Tsai --- .../scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala | 1 + 1 file changed, 1 insertion(+) diff --git a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala index 5cbb0654eb376..16522e981a338 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -5827,6 +5827,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti // Map stage completes successfully, completeShuffleMapStageSuccessfully(0, 0, 3, Seq("hostA", "hostB")) taskIdCount += 2 + sc.listenerBus.waitUntilEmpty() assert(completedStage === List(0)) // Now submit the first reducer stage