refine UT

core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala

@@ -401,14 +401,14 @@ class TaskSetManagerSuite extends FunSuite with LocalSparkContext with Logging {
     // All tasks added to no-pref list since no preferred location is available
     assert(manager.pendingTasksWithNoPrefs.size === 4)
     // Only ANY is valid
-    assert(manager.myLocalityLevels.size === 1)
+    assert(manager.myLocalityLevels.sameElements(Array(ANY)))
     // Add a new executor
     sched.addExecutor(("execD", "host1"))
     manager.executorAdded()
     // Task 0 and 1 should be removed from no-pref list
     assert(manager.pendingTasksWithNoPrefs.size === 2)
     // Valid locality should contain NODE_LOCAL and ANY
-    assert(manager.myLocalityLevels.size === 2)
+    assert(manager.myLocalityLevels.sameElements(Array(NODE_LOCAL, ANY)))
     // Offer host1, execD, at PROCESS_LOCAL level: task 0 should be chosen
     // because PROCESS_LOCAL is not valid at the moment
     assert(manager.resourceOffer("execD", "host1", PROCESS_LOCAL).get.index === 0)
@@ -418,7 +418,7 @@ class TaskSetManagerSuite extends FunSuite with LocalSparkContext with Logging {
     // No-pref list now only contains task 3
     assert(manager.pendingTasksWithNoPrefs.size === 1)
     // Valid locality should contain PROCESS_LOCAL, NODE_LOCAL and ANY
-    assert(manager.myLocalityLevels.size === 3)
+    assert(manager.myLocalityLevels.sameElements(Array(PROCESS_LOCAL, NODE_LOCAL, ANY)))
     // Offer host2, execC, at PROCESS_LOCAL level: task 2 should be chosen
     assert(manager.resourceOffer("execC", "host2", PROCESS_LOCAL).get.index === 2)
     // Offer host1, execD again at PROCESS_LOCAL level: task 3 should be chosen