From 063142e4f594308f2e321855672bcb0d8ba9138d Mon Sep 17 00:00:00 2001
From: jiang13021 <jiangyanze.jyz@antgroup.com>
Date: Sat, 16 Nov 2024 11:18:57 +0800
Subject: [PATCH] fix ut in spark-2.4

---
 .../deploy/cluster/PushMergedDataHardSplitSuite.scala  | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/worker/src/test/scala/org/apache/celeborn/service/deploy/cluster/PushMergedDataHardSplitSuite.scala b/worker/src/test/scala/org/apache/celeborn/service/deploy/cluster/PushMergedDataHardSplitSuite.scala
index b09b529f2e..3d1b76a667 100644
--- a/worker/src/test/scala/org/apache/celeborn/service/deploy/cluster/PushMergedDataHardSplitSuite.scala
+++ b/worker/src/test/scala/org/apache/celeborn/service/deploy/cluster/PushMergedDataHardSplitSuite.scala
@@ -88,11 +88,11 @@ class PushMergedDataHardSplitSuite extends AnyFunSuite
     val partitionLocationMap =
       shuffleClient.getPartitionLocation(SHUFFLE_ID, MAP_NUM, PARTITION_NUM)
     val worker2PartitionIds = mutable.Map.empty[WorkerInfo, ArrayBuffer[Int]]
-    partitionLocationMap.forEach {
-      (partitionId, partitionLocation) =>
-        worker2PartitionIds
-          .getOrElseUpdate(partitionLocation.getWorker, ArrayBuffer.empty)
-          .append(partitionId)
+    for (partitionId <- 0 until PARTITION_NUM) {
+      val partitionLocation = partitionLocationMap.get(partitionId)
+      worker2PartitionIds
+        .getOrElseUpdate(partitionLocation.getWorker, ArrayBuffer.empty)
+        .append(partitionId)
     }
     val partitions = worker2PartitionIds.values.filter(_.size >= 2).head
     assert(partitions.length >= 2)