[SPARK-15354][FLAKY-TEST] TopologyAwareBlockReplicationPolicyBehavior…

….Peers in 2 racks ## What changes were proposed in this pull request? `TopologyAwareBlockReplicationPolicyBehavior.Peers in 2 racks` is failing occasionally: https://spark-tests.appspot.com/test-details?suite_name=org.apache.spark.storage.TopologyAwareBlockReplicationPolicyBehavior&test_name=Peers+in+2+racks. This is because, when we generate 10 block manager id to test, they may all belong to the same rack, as the rack is randomly picked. This PR fixes this problem by forcing each rack to be picked at least once. ## How was this patch tested? N/A Author: Wenchen Fan <[email protected]> Closes apache#17624 from cloud-fan/test.
dbtsai · Apr 13, 2017 · a7b430b · a7b430b
1 parent 924c424
commit a7b430b
Showing 1 changed file with 11 additions and 2 deletions.
diff --git a/core/src/test/scala/org/apache/spark/storage/BlockReplicationPolicySuite.scala b/core/src/test/scala/org/apache/spark/storage/BlockReplicationPolicySuite.scala
@@ -70,9 +70,18 @@ class RandomBlockReplicationPolicyBehavior extends SparkFunSuite
     }
   }
 
+  /**
+   * Returns a sequence of [[BlockManagerId]], whose rack is randomly picked from the given `racks`.
+   * Note that, each rack will be picked at least once from `racks`, if `count` is greater or equal
+   * to the number of `racks`.
+   */
   protected def generateBlockManagerIds(count: Int, racks: Seq[String]): Seq[BlockManagerId] = {
-    (1 to count).map{i =>
-      BlockManagerId(s"Exec-$i", s"Host-$i", 10000 + i, Some(racks(Random.nextInt(racks.size))))
+    val randomizedRacks: Seq[String] = Random.shuffle(
+      racks ++ racks.length.until(count).map(_ => racks(Random.nextInt(racks.length)))
+    )
+
+    (0 until count).map { i =>
+      BlockManagerId(s"Exec-$i", s"Host-$i", 10000 + i, Some(randomizedRacks(i)))
     }
   }
 }