Spark unit test context

OptimizerSpec.scala

class OptimizerSpec extends FlatSpec with Matchers with SparkSetup {
    "solve" should "work with just intercept" in withSparkSession { session =>
...
   }
}

SparkSetup.scala

trait SparkSetup {
  private def ignore(x: Any): Unit = {}

  private val conf: SparkConf = new SparkConf()
    .setMaster("local[1]")
    .setAppName("Spark test")
    .set("spark.sql.shuffle.partitions", "1")

  def withSparkContext(testMethod: (SparkContext) => Any): Unit = {
    val sparkContext = new SparkContext(conf)
    try {
      ignore(testMethod(sparkContext))
    }
    finally sparkContext.stop()
  }

  def withSparkSession(testMethod: (SparkSession) => Any): Unit = {
    val sparkSession = SparkSession.builder().config(conf).getOrCreate()
    try {
      ignore(testMethod(sparkSession))
    }
    finally sparkSession.stop()
  }
}