diff --git a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/PushDataTimeoutTest.scala b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/PushDataTimeoutTest.scala
index 5bf7c1303b1..ea398968cd2 100644
--- a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/PushDataTimeoutTest.scala
+++ b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/PushDataTimeoutTest.scala
@@ -144,9 +144,12 @@ class PushDataTimeoutTest extends AnyFunSuite
       .getLifecycleManager
       .workerStatusTracker
       .excludedWorkers
+      .asScala.filter { case (_, (code, _)) =>
+        code != StatusCode.WORKER_UNKNOWN
+      }.toMap
 
-    assert(excludedWorkers.size() > 0)
-    excludedWorkers.asScala.foreach { case (_, (code, _)) =>
+    assert(excludedWorkers.size > 0)
+    excludedWorkers.foreach { case (_, (code, _)) =>
       assert(code == StatusCode.PUSH_DATA_TIMEOUT_PRIMARY ||
         code == StatusCode.PUSH_DATA_TIMEOUT_REPLICA)
     }