Minimize sql.SparkSessionBinder stuff

Author: Matthis Gördel

sql/core/src/test/scala/org/apache/spark/sql/SparkSessionBinder.scala

@@ -19,17 +19,17 @@ package org.apache.spark.sql
 
 import scala.concurrent.duration._
 
-import org.scalatest.{BeforeAndAfterEach, Suite}
+import org.scalatest.{BeforeAndAfterAll, BeforeAndAfterEach, Suite}
 import org.scalatest.concurrent.Eventually
 
-import org.apache.spark.{DebugFilesystem, SparkConf}
+import org.apache.spark.{DebugFilesystem, SparkConf, SparkFunSuite}
 import org.apache.spark.internal.config.UNSAFE_EXCEPTION_ON_MEMORY_LEAK
 import org.apache.spark.sql.catalyst.expressions.CodegenObjectFactoryMode
 import org.apache.spark.sql.catalyst.optimizer.ConvertToLocalRelation
 import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf}
 import org.apache.spark.sql.test.TestSparkSession
 
-trait SparkSessionBinder extends QueryTest with SparkSessionBinderBase {
+trait SparkSessionBinder extends SparkSessionBinderBase { self: SparkFunSuite =>
 
   /**
    * Suites extending this trait are sharing resources (e.g. SparkSession) in their
@@ -56,9 +56,9 @@ trait SparkSessionBinder extends QueryTest with SparkSessionBinderBase {
 }
 
 trait SparkSessionBinderBase
-  extends QueryTestBase
-  with SparkSessionProvider
+  extends SparkSessionProvider
   with BeforeAndAfterEach
+  with BeforeAndAfterAll
   with Eventually { self: Suite =>
 
   protected def sparkConf = {

sql/core/src/test/scala/org/apache/spark/sql/classic/SparkSessionBinder.scala

@@ -17,40 +17,11 @@
 
 package org.apache.spark.sql.classic
 
-import scala.concurrent.duration._
-
-import org.apache.spark.sql
+import org.apache.spark.{sql, SparkFunSuite}
 
 /**
  * Overrides [[spark]] to provide a [[SparkSession classic.SparkSession]]
  */
-trait SparkSessionBinder extends sql.SparkSessionBinder with QueryTest {
+trait SparkSessionBinder extends sql.SparkSessionBinder { self: SparkFunSuite =>
   override protected def spark: SparkSession = super.spark.asInstanceOf[SparkSession]
-
-  // Runs func (which must trigger exactly one SQL execution) and returns the SQL metrics of that
-  // execution as a map keyed by (planNodeId, planNodeName, metricName) -> metricValue.
-  def runAndFetchMetrics(func: => Unit): Map[(Long, String, String), String] = {
-    val statusStore = spark.sharedState.statusStore
-    val oldCount = statusStore.executionsList().size
-
-    func
-
-    // Wait until the new execution is started and being tracked.
-    eventually(timeout(10.seconds), interval(10.milliseconds)) {
-      assert(statusStore.executionsCount() >= oldCount)
-    }
-
-    // Wait for listener to finish computing the metrics for the execution.
-    eventually(timeout(10.seconds), interval(10.milliseconds)) {
-      assert(statusStore.executionsList().nonEmpty &&
-        statusStore.executionsList().last.metricValues != null)
-    }
-
-    val exec = statusStore.executionsList().last
-    val execId = exec.executionId
-    val sqlMetrics = statusStore.planGraph(execId).allNodes
-      .flatMap(n => n.metrics.map(m => (m.accumulatorId, (n.id, n.name, m.name))))
-      .toMap
-    statusStore.executionMetrics(execId).map { case (k, v) => sqlMetrics(k) -> v }
-  }
 }

sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala

@@ -17,20 +17,50 @@
 
 package org.apache.spark.sql.test
 
+import scala.concurrent.duration._
+
 import org.scalatest.Suite
 
-import org.apache.spark.sql
+import org.apache.spark.sql.{QueryTestBase, QueryTest, SparkSessionBinderBase}
 import org.apache.spark.sql.classic
 
 @deprecated("Use SparkSessionBinder (or classic.SparkSessionBinder if required) instead", "4.2.0")
-trait SharedSparkSession extends classic.SparkSessionBinder
+trait SharedSparkSession extends QueryTest with classic.SparkSessionBinder {
+
+  // Runs func (which must trigger exactly one SQL execution) and returns the SQL metrics of that
+  // execution as a map keyed by (planNodeId, planNodeName, metricName) -> metricValue.
+  def runAndFetchMetrics(func: => Unit): Map[(Long, String, String), String] = {
+    val statusStore = spark.sharedState.statusStore
+    val oldCount = statusStore.executionsList().size
+
+    func
+
+    // Wait until the new execution is started and being tracked.
+    eventually(timeout(10.seconds), interval(10.milliseconds)) {
+      assert(statusStore.executionsCount() >= oldCount)
+    }
+
+    // Wait for listener to finish computing the metrics for the execution.
+    eventually(timeout(10.seconds), interval(10.milliseconds)) {
+      assert(statusStore.executionsList().nonEmpty &&
+        statusStore.executionsList().last.metricValues != null)
+    }
+
+    val exec = statusStore.executionsList().last
+    val execId = exec.executionId
+    val sqlMetrics = statusStore.planGraph(execId).allNodes
+      .flatMap(n => n.metrics.map(m => (m.accumulatorId, (n.id, n.name, m.name))))
+      .toMap
+    statusStore.executionMetrics(execId).map { case (k, v) => sqlMetrics(k) -> v }
+  }
+}
 
 
 /**
  * Helper trait for SQL test suites where all tests share a single [[TestSparkSession]].
  */
 @deprecated("Use SparkSessionBinder (or classic.SparkSessionBinder if required) instead", "4.2.0")
-trait SharedSparkSessionBase extends sql.SparkSessionBinderBase { self: Suite =>
+trait SharedSparkSessionBase extends QueryTestBase with SparkSessionBinderBase { self: Suite =>
 
   protected override def spark: classic.SparkSession =
     super.spark.asInstanceOf[classic.SparkSession]