Move all SparkFunSuite-dependent methods from SQLTestUtils to QueryTest

zhengruifeng · zhengruifeng · commit 33167c4170b0 · 2026-04-09T11:25:35.000Z
SQLTestUtils body is now empty. Methods like withTempDir, testQuietly,
testWithWholeStageCodegenOnAndOff are in QueryTest. SparkPlanTest,
PlannerSuite, and HDFSMetadataLogSuite updated to extend QueryTest.

Co-authored-by: Isaac
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/QueryTest.scala b/sql/core/src/test/scala/org/apache/spark/sql/QueryTest.scala
@@ -28,7 +28,8 @@ import scala.jdk.CollectionConverters._
 import scala.language.implicitConversions
 
 import org.apache.hadoop.fs.Path
-import org.scalatest.{Assertions, BeforeAndAfterAll, Suite}
+import org.scalactic.source.Position
+import org.scalatest.{Assertions, BeforeAndAfterAll, Suite, Tag}
 import org.scalatest.concurrent.Eventually
 
 import org.apache.spark.SparkFunSuite
@@ -43,12 +44,14 @@ import org.apache.spark.sql.catalyst.util._
 import org.apache.spark.sql.classic.{ClassicConversions, ColumnConversions, ColumnNodeToExpressionConverter, SQLImplicits}
 import org.apache.spark.sql.classic.ClassicConversions._
 import org.apache.spark.sql.execution.{FilterExec, QueryExecution, SQLExecution}
+import org.apache.spark.sql.execution.adaptive.DisableAdaptiveExecution
 import org.apache.spark.sql.execution.columnar.InMemoryRelation
 import org.apache.spark.sql.execution.datasources.DataSourceUtils
+import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.util.QueryExecutionListener
 import org.apache.spark.storage.StorageLevel
+import org.apache.spark.util.{UninterruptibleThread, Utils}
 import org.apache.spark.util.ArrayImplicits._
-import org.apache.spark.util.Utils
 
 
 /**
@@ -602,7 +605,118 @@ trait QueryTestBase extends PlanTestBase
 
 }
 
-abstract class QueryTest extends SparkFunSuite with QueryTestBase with PlanTest
+abstract class QueryTest extends SparkFunSuite with QueryTestBase with PlanTest {
+  import scala.util.control.NonFatal
+
+  /**
+   * Creates a temporary directory, which is then passed to `f` and will be deleted after `f`
+   * returns.
+   */
+  protected override def withTempDir(f: File => Unit): Unit = {
+    super.withTempDir { dir =>
+      f(dir)
+      waitForTasksToFinish()
+    }
+  }
+
+  /**
+   * A helper function for turning off/on codegen.
+   */
+  protected def testWithWholeStageCodegenOnAndOff(testName: String)(f: String => Unit): Unit = {
+    Seq("false", "true").foreach { codegenEnabled =>
+      val isTurnOn = if (codegenEnabled == "true") "on" else "off"
+      test(s"$testName (whole-stage-codegen ${isTurnOn})") {
+        withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> codegenEnabled) {
+          f(codegenEnabled)
+        }
+      }
+    }
+  }
+
+  /**
+   * Disable stdout and stderr when running the test. To not output the logs to the console,
+   * ConsoleAppender's `follow` should be set to `true` so that it will honor reassignments of
+   * System.out or System.err. Otherwise, ConsoleAppender will still output to the console even if
+   * we change System.out and System.err.
+   */
+  protected def testQuietly(name: String)(f: => Unit): Unit = {
+    test(name) {
+      quietly {
+        f
+      }
+    }
+  }
+
+  override protected def test(testName: String, testTags: Tag*)(testFun: => Any)
+      (implicit pos: Position): Unit = {
+    if (testTags.exists(_.isInstanceOf[DisableAdaptiveExecution])) {
+      super.test(testName, testTags: _*) {
+        withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+          testFun
+        }
+      }
+    } else {
+      super.test(testName, testTags: _*)(testFun)
+    }
+  }
+
+  /**
+   * Run a test on a separate `UninterruptibleThread`.
+   */
+  protected def testWithUninterruptibleThread(name: String, quietly: Boolean = false)
+    (body: => Unit): Unit = {
+    val timeoutMillis = 10000
+    @transient var ex: Throwable = null
+
+    def runOnThread(): Unit = {
+      val thread = new UninterruptibleThread(s"Testing thread for test $name") {
+        override def run(): Unit = {
+          try {
+            body
+          } catch {
+            case NonFatal(e) =>
+              ex = e
+          }
+        }
+      }
+      thread.setDaemon(true)
+      thread.start()
+      thread.join(timeoutMillis)
+      if (thread.isAlive) {
+        thread.interrupt()
+        // If this interrupt does not work, then this thread is most likely running something that
+        // is not interruptible. There is not much point to wait for the thread to terminate, and
+        // we rather let the JVM terminate the thread on exit.
+        fail(
+          s"Test '$name' running on o.a.s.util.UninterruptibleThread timed out after" +
+            s" $timeoutMillis ms")
+      } else if (ex != null) {
+        throw ex
+      }
+    }
+
+    if (quietly) {
+      testQuietly(name) { runOnThread() }
+    } else {
+      test(name) { runOnThread() }
+    }
+  }
+
+  /**
+   * Copy file in jar's resource to a temp file, then pass it to `f`.
+   * This function is used to make `f` can use the path of temp file(e.g. file:/), instead of
+   * path of jar's resource which starts with 'jar:file:/'
+   */
+  protected def withResourceTempPath(resourcePath: String)(f: File => Unit): Unit = {
+    val inputStream =
+      Thread.currentThread().getContextClassLoader.getResourceAsStream(resourcePath)
+    withTempDir { dir =>
+      val tmpFile = new File(dir, "tmp")
+      Files.copy(inputStream, tmpFile.toPath)
+      f(tmpFile)
+    }
+  }
+}
 
 object QueryTest extends Assertions {
   /**
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala
@@ -19,7 +19,7 @@ package org.apache.spark.sql.execution
 
 import org.apache.spark.SparkUnsupportedOperationException
 import org.apache.spark.rdd.RDD
-import org.apache.spark.sql.{execution, DataFrame, Row}
+import org.apache.spark.sql.{execution, DataFrame, QueryTest, Row}
 import org.apache.spark.sql.AnalysisException
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions._
@@ -37,7 +37,7 @@ import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.test.SharedSparkSession
 import org.apache.spark.sql.types._
 
-class PlannerSuite extends SharedSparkSession with AdaptiveSparkPlanHelper {
+class PlannerSuite extends QueryTest with SharedSparkSession with AdaptiveSparkPlanHelper {
   import testImplicits._
 
   setupTestData()
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanTest.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanTest.scala
@@ -19,8 +19,7 @@ package org.apache.spark.sql.execution
 
 import scala.util.control.NonFatal
 
-import org.apache.spark.SparkFunSuite
-import org.apache.spark.sql.{classic, DataFrame, Row, SparkSessionProvider, SQLContext}
+import org.apache.spark.sql.{classic, DataFrame, QueryTest, Row, SQLContext}
 import org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute
 import org.apache.spark.sql.catalyst.plans.logical.LocalRelation
 import org.apache.spark.sql.classic.ClassicConversions._
@@ -30,7 +29,7 @@ import org.apache.spark.sql.test.SQLTestUtils
  * Base class for writing tests for individual physical operators. For an example of how this
  * class's test helper methods can be used, see [[SortSuite]].
  */
-private[sql] abstract class SparkPlanTest extends SparkFunSuite with SparkSessionProvider {
+private[sql] abstract class SparkPlanTest extends QueryTest {
   override protected def spark: classic.SparkSession
 
   /**
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/HDFSMetadataLogSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/HDFSMetadataLogSuite.scala
@@ -25,12 +25,13 @@ import scala.language.implicitConversions
 import org.scalatest.concurrent.Waiters._
 import org.scalatest.time.SpanSugar._
 
+import org.apache.spark.sql.QueryTest
 import org.apache.spark.sql.execution.streaming.checkpointing.{FileContextBasedCheckpointFileManager, FileSystemBasedCheckpointFileManager, HDFSMetadataLog}
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.test.SharedSparkSession
 import org.apache.spark.util.UninterruptibleThread
 
-class HDFSMetadataLogSuite extends SharedSparkSession {
+class HDFSMetadataLogSuite extends QueryTest with SharedSparkSession {
 
   private implicit def toOption[A](a: A): Option[A] = Option(a)
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestUtils.scala b/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestUtils.scala
@@ -17,21 +17,11 @@
 
 package org.apache.spark.sql.test
 
-import java.io.File
-import java.nio.file.Files
-
-import scala.util.control.NonFatal
-
-import org.scalactic.source.Position
-import org.scalatest.{Suite, Tag}
+import org.scalatest.Suite
 
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.sql.{QueryTest, QueryTestBase, Row}
 import org.apache.spark.sql.catalyst.plans.PlanTest
-import org.apache.spark.sql.catalyst.util.quietly
-import org.apache.spark.sql.execution.adaptive.DisableAdaptiveExecution
-import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.util.UninterruptibleThread
 
 /**
  * Helper trait that can be extended by all external SQL test suites.
@@ -44,120 +34,9 @@ private[sql] trait SQLTestUtilsBase extends QueryTestBase { self: Suite => }
  * Helper trait that should be extended by all SQL test suites within the Spark
  * code base.
  *
- * This extends [[SparkFunSuite]] with [[QueryTestBase]] and adds test helpers that
- * depend on [[SparkFunSuite]] (e.g. `test()` overrides, `withTempDir`).
+ * This is now an alias for [[QueryTest]].
  */
-private[sql] trait SQLTestUtils extends SparkFunSuite with SQLTestUtilsBase with PlanTest {
-
-  /**
-   * Creates a temporary directory, which is then passed to `f` and will be deleted after `f`
-   * returns.
-   */
-  protected override def withTempDir(f: File => Unit): Unit = {
-    super.withTempDir { dir =>
-      f(dir)
-      waitForTasksToFinish()
-    }
-  }
-
-  /**
-   * A helper function for turning off/on codegen.
-   */
-  protected def testWithWholeStageCodegenOnAndOff(testName: String)(f: String => Unit): Unit = {
-    Seq("false", "true").foreach { codegenEnabled =>
-      val isTurnOn = if (codegenEnabled == "true") "on" else "off"
-      test(s"$testName (whole-stage-codegen ${isTurnOn})") {
-        withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> codegenEnabled) {
-          f(codegenEnabled)
-        }
-      }
-    }
-  }
-
-  /**
-   * Disable stdout and stderr when running the test. To not output the logs to the console,
-   * ConsoleAppender's `follow` should be set to `true` so that it will honor reassignments of
-   * System.out or System.err. Otherwise, ConsoleAppender will still output to the console even if
-   * we change System.out and System.err.
-   */
-  protected def testQuietly(name: String)(f: => Unit): Unit = {
-    test(name) {
-      quietly {
-        f
-      }
-    }
-  }
-
-  override protected def test(testName: String, testTags: Tag*)(testFun: => Any)
-      (implicit pos: Position): Unit = {
-    if (testTags.exists(_.isInstanceOf[DisableAdaptiveExecution])) {
-      super.test(testName, testTags: _*) {
-        withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
-          testFun
-        }
-      }
-    } else {
-      super.test(testName, testTags: _*)(testFun)
-    }
-  }
-
-  /**
-   * Run a test on a separate `UninterruptibleThread`.
-   */
-  protected def testWithUninterruptibleThread(name: String, quietly: Boolean = false)
-    (body: => Unit): Unit = {
-    val timeoutMillis = 10000
-    @transient var ex: Throwable = null
-
-    def runOnThread(): Unit = {
-      val thread = new UninterruptibleThread(s"Testing thread for test $name") {
-        override def run(): Unit = {
-          try {
-            body
-          } catch {
-            case NonFatal(e) =>
-              ex = e
-          }
-        }
-      }
-      thread.setDaemon(true)
-      thread.start()
-      thread.join(timeoutMillis)
-      if (thread.isAlive) {
-        thread.interrupt()
-        // If this interrupt does not work, then this thread is most likely running something that
-        // is not interruptible. There is not much point to wait for the thread to terminate, and
-        // we rather let the JVM terminate the thread on exit.
-        fail(
-          s"Test '$name' running on o.a.s.util.UninterruptibleThread timed out after" +
-            s" $timeoutMillis ms")
-      } else if (ex != null) {
-        throw ex
-      }
-    }
-
-    if (quietly) {
-      testQuietly(name) { runOnThread() }
-    } else {
-      test(name) { runOnThread() }
-    }
-  }
-
-  /**
-   * Copy file in jar's resource to a temp file, then pass it to `f`.
-   * This function is used to make `f` can use the path of temp file(e.g. file:/), instead of
-   * path of jar's resource which starts with 'jar:file:/'
-   */
-  protected def withResourceTempPath(resourcePath: String)(f: File => Unit): Unit = {
-    val inputStream =
-      Thread.currentThread().getContextClassLoader.getResourceAsStream(resourcePath)
-    withTempDir { dir =>
-      val tmpFile = new File(dir, "tmp")
-      Files.copy(inputStream, tmpFile.toPath)
-      f(tmpFile)
-    }
-  }
-}
+private[sql] trait SQLTestUtils extends SparkFunSuite with SQLTestUtilsBase with PlanTest
 
 private[sql] object SQLTestUtils {
 
