Skip to content
21 changes: 4 additions & 17 deletions core/src/test/scala/org/apache/spark/SparkTestSuite.scala
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ package org.apache.spark

import java.io.File
import java.nio.charset.StandardCharsets.UTF_8
import java.nio.file.{Files, Path}
import java.nio.file.Files
import java.util.{Locale, TimeZone}

import scala.collection.mutable.{ArrayBuffer, ListBuffer}
Expand All @@ -36,7 +36,7 @@ import org.scalatest.concurrent.TimeLimits
import org.apache.spark.deploy.LocalSparkCluster
import org.apache.spark.internal.Logging
import org.apache.spark.internal.config.Tests.IS_TESTING
import org.apache.spark.util.{AccumulatorContext, Utils}
import org.apache.spark.util.{AccumulatorContext, TestEnvHelper, Utils}

/**
* Base trait for all unit tests in Spark for handling common functionality.
Expand Down Expand Up @@ -70,7 +70,8 @@ trait SparkTestSuite
with BeforeAndAfterEach
with ThreadAudit
with TimeLimits
with Logging {
with Logging
with TestEnvHelper {
// scalastyle:on

// Initialize the logger forcibly to let the logger log timestamp
Expand All @@ -88,9 +89,6 @@ trait SparkTestSuite

protected def enableAutoThreadAudit = true

protected def regenerateGoldenFiles: Boolean =
System.getenv("SPARK_GENERATE_GOLDEN_FILES") == "1"

protected override def beforeAll(): Unit = {
System.setProperty(IS_TESTING.key, "true")
if (enableAutoThreadAudit) {
Expand Down Expand Up @@ -131,17 +129,6 @@ trait SparkTestSuite
file
}

/**
* Get a Path relative to the root project. It is assumed that a spark home is set.
*/
protected final def getWorkspaceFilePath(first: String, more: String*): Path = {
if (!(sys.props.contains("spark.test.home") || sys.env.contains("SPARK_HOME"))) {
fail("spark.test.home or SPARK_HOME is not set.")
}
val sparkHome = sys.props.getOrElse("spark.test.home", sys.env("SPARK_HOME"))
java.nio.file.Paths.get(sparkHome, first +: more: _*)
}

// subclasses can override this to exclude certain tests by name
// useful when inheriting a test suite but do not want to run all tests in it
protected def excluded: Seq[String] = Seq.empty
Expand Down
46 changes: 46 additions & 0 deletions core/src/test/scala/org/apache/spark/util/TestEnvHelper.scala
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

package org.apache.spark.util

import java.nio.file.{Path, Paths}

import org.scalatest.Assertions

/**
* Shared test helpers for resolving workspace files and reading golden-file env-var flags.
* Mixed into `SparkTestSuite` so every suite extending `SparkFunSuite` inherits these.
*/
trait TestEnvHelper extends Assertions {

/**
* Get a Path relative to the root project. It is assumed that a spark home is set.
*/
protected final def getWorkspaceFilePath(first: String, more: String*): Path = {
if (!(sys.props.contains("spark.test.home") || sys.env.contains("SPARK_HOME"))) {
fail("spark.test.home or SPARK_HOME is not set.")
}
val sparkHome = sys.props.getOrElse("spark.test.home", sys.env("SPARK_HOME"))
Paths.get(sparkHome, first +: more: _*)
}

protected def regenerateGoldenFiles: Boolean =
System.getenv("SPARK_GENERATE_GOLDEN_FILES") == "1"

protected def cleanOrphanedGoldenFiles: Boolean =
System.getenv("SPARK_CLEAN_ORPHANED_GOLDEN_FILES") == "1"
}
Original file line number Diff line number Diff line change
Expand Up @@ -75,9 +75,6 @@ class ProtoToParsedPlanTestSuite
with SharedSparkSession
with ResourceHelper {

private val cleanOrphanedGoldenFiles: Boolean =
System.getenv("SPARK_CLEAN_ORPHANED_GOLDEN_FILES") == "1"

val url = "jdbc:h2:mem:testdb0"
var conn: java.sql.Connection = null

Expand Down