[GitHub] attilapiros commented on a change in pull request #23801: [SPARK-26891][YARN] Fixing flaky test in YarnSchedulerBackendSuite

[GitBox <gi...@apache.org>]

attilapiros commented on a change in pull request #23801: [SPARK-26891][YARN] Fixing flaky test in YarnSchedulerBackendSuite
URL: https://github.com/apache/spark/pull/23801#discussion_r257517645
 
 

 ##########
 File path: resource-managers/yarn/src/test/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackendSuite.scala
 ##########
 @@ -32,15 +33,33 @@ import org.apache.spark.ui.TestFilter
 
 class YarnSchedulerBackendSuite extends SparkFunSuite with MockitoSugar with LocalSparkContext {
 
+  private var yarnSchedulerBackend: YarnSchedulerBackend = _
+
+  override def afterEach() {
+    try {
+      yarnSchedulerBackend.stop()
+    } finally {
+      super.afterEach()
+    }
+  }
+
   test("RequestExecutors reflects node blacklist and is serializable") {
     sc = new SparkContext("local", "YarnSchedulerBackendSuite")
-    val sched = mock[TaskSchedulerImpl]
-    when(sched.sc).thenReturn(sc)
-    val yarnSchedulerBackend = new YarnSchedulerBackend(sched, sc) {
+    // Subclassing the TaskSchedulerImpl here instead of using Mockito. For details see SPARK-26891.
+    val sched = new TaskSchedulerImpl(sc) {
+      val blacklistedNodes = new AtomicReference[Set[String]]()
+
+      def setNodeBlacklist(nodeBlacklist: Set[String]): Unit = blacklistedNodes.set(nodeBlacklist)
+
+      override def nodeBlacklist(): Set[String] = blacklistedNodes.get()
+    }
+
+    val yarnSchedulerBackendExtended = new YarnSchedulerBackend(sched, sc) {
 
 Review comment:
   It is needed because of the different type: the `yarnSchedulerBackend` type is `YarnSchedulerBackend` but `yarnSchedulerBackendExtended` type is an anonim subclass of `YarnSchedulerBackend` with the extra def `setNodeBlacklist`. On `yarnSchedulerBackend` I cannot call this extra method.
   

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services

---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org