Avoid the need to switch to HiveContexts

This is a clean up to refactor helper test traits and abstract
classes in such a way that is accessible to hive tests.

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

@@ -30,7 +30,7 @@ import org.apache.spark.sql.test.MyTestSQLContext
 private case class BigData(s: String)
 
 class CachedTableSuite extends QueryTest with MyTestSQLContext {
-  private val ctx = sqlContextWithData
+  private val ctx = sqlContext
   import ctx.implicits._
   import ctx._
 

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

@@ -25,7 +25,7 @@ import org.apache.spark.sql.types._
 import org.apache.spark.sql.test.SQLTestUtils
 
 class ColumnExpressionSuite extends QueryTest with SQLTestUtils {
-  private val ctx = sqlContextWithData
+  private val ctx = sqlContext
   import ctx.implicits._
   import ctx._
 

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

@@ -23,7 +23,7 @@ import org.apache.spark.sql.types.DecimalType
 
 
 class DataFrameAggregateSuite extends QueryTest with MyTestSQLContext {
-  private val ctx = sqlContextWithData
+  private val ctx = sqlContext
   import ctx.implicits._
   import ctx._
 

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

@@ -25,7 +25,7 @@ import org.apache.spark.sql.types._
  * Test suite for functions in [[org.apache.spark.sql.functions]].
  */
 class DataFrameFunctionsSuite extends QueryTest with MyTestSQLContext {
-  private val ctx = sqlContextWithData
+  private val ctx = sqlContext
   import ctx.implicits._
   import ctx._
 

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

@@ -22,7 +22,7 @@ import org.apache.spark.sql.functions._
 import org.apache.spark.sql.test.MyTestSQLContext
 
 class DataFrameJoinSuite extends QueryTest with MyTestSQLContext {
-  private val ctx = sqlContextWithData
+  private val ctx = sqlContext
   import ctx.implicits._
   import ctx._
 

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

@@ -31,7 +31,7 @@ import org.apache.spark.sql.types._
 import org.apache.spark.sql.test.{ExamplePointUDT, ExamplePoint, SQLTestUtils}
 
 class DataFrameSuite extends QueryTest with SQLTestUtils {
-  private val ctx = sqlContextWithData
+  private val ctx = sqlContext
   import ctx.implicits._
   import ctx._
 

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

@@ -25,7 +25,7 @@ import org.apache.spark.sql.test.MyTestSQLContext
 
 
 class JoinSuite extends QueryTest with BeforeAndAfterEach with MyTestSQLContext {
-  private val ctx = sqlContextWithData
+  private val ctx = sqlContext
   import ctx.implicits._
   import ctx._
 

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

@@ -33,7 +33,7 @@ import org.apache.spark.sql.types._
 class MyDialect extends DefaultParserDialect
 
 class SQLQuerySuite extends QueryTest with SQLTestUtils {
-  private val ctx = sqlContextWithData
+  private val ctx = sqlContext
   import ctx.implicits._
   import ctx._
 

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

@@ -22,7 +22,7 @@ import org.apache.spark.sql.test.SQLTestUtils
 private case class FunctionResult(f1: String, f2: String)
 
 class UDFSuite extends QueryTest with SQLTestUtils {
-  private val ctx = sqlContextWithData
+  private val ctx = sqlContext
   import ctx.implicits._
   import ctx._
 

sql/core/src/test/scala/org/apache/spark/sql/columnar/InMemoryColumnarQuerySuite.scala

@@ -25,7 +25,7 @@ import org.apache.spark.sql.{QueryTest, Row}
 import org.apache.spark.storage.StorageLevel.MEMORY_ONLY
 
 class InMemoryColumnarQuerySuite extends QueryTest with MyTestSQLContext {
-  private val ctx = sqlContextWithData
+  private val ctx = sqlContext
   import ctx.implicits._
   import ctx._
 

sql/core/src/test/scala/org/apache/spark/sql/columnar/PartitionBatchPruningSuite.scala

@@ -24,7 +24,7 @@ import org.apache.spark.sql._
 import org.apache.spark.sql.test.MyTestSQLContext
 
 class PartitionBatchPruningSuite extends SparkFunSuite with BeforeAndAfter with MyTestSQLContext {
-  private val ctx = sqlContextWithData
+  private val ctx = sqlContext
   import ctx.implicits._
   import ctx._
 

sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala

@@ -28,7 +28,7 @@ import org.apache.spark.sql.{execution, Row, SQLConf}
 
 
 class PlannerSuite extends SparkFunSuite with SQLTestUtils {
-  private val ctx = sqlContextWithData
+  private val ctx = sqlContext
   import ctx.implicits._
   import ctx.planner._
   import ctx._

sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanTest.scala

@@ -27,18 +27,25 @@ import org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute
 import org.apache.spark.sql.catalyst.util._
 import org.apache.spark.sql.test.MyTestSQLContext
 
-
 /**
  * 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]].
  */
-abstract class SparkPlanTest extends SparkFunSuite with MyTestSQLContext {
+private[sql] abstract class SparkPlanTest extends AbstractSparkPlanTest with MyTestSQLContext {
+  protected override def _sqlContext: SQLContext = sqlContext
+}
+
+/**
+ * Helper class for testing individual physical operators with a pluggable [[SQLContext]].
+ */
+private[sql] abstract class AbstractSparkPlanTest extends SparkFunSuite {
+  protected def _sqlContext: SQLContext
 
   /**
    * Creates a DataFrame from a local Seq of Product.
    */
   implicit def localSeqToDataFrameHolder[A <: Product : TypeTag](data: Seq[A]): DataFrameHolder = {
-    sqlContext.implicits.localSeqToDataFrameHolder(data)
+    _sqlContext.implicits.localSeqToDataFrameHolder(data)
   }
 
   /**
@@ -99,7 +106,7 @@ abstract class SparkPlanTest extends SparkFunSuite with MyTestSQLContext {
       planFunction: Seq[SparkPlan] => SparkPlan,
       expectedAnswer: Seq[Row],
       sortAnswers: Boolean = true): Unit = {
-    SparkPlanTest.checkAnswer(input, planFunction, expectedAnswer, sortAnswers, sqlContext) match {
+    SparkPlanTest.checkAnswer(input, planFunction, expectedAnswer, sortAnswers, _sqlContext) match {
       case Some(errorMessage) => fail(errorMessage)
       case None =>
     }
@@ -123,7 +130,7 @@ abstract class SparkPlanTest extends SparkFunSuite with MyTestSQLContext {
       expectedPlanFunction: SparkPlan => SparkPlan,
       sortAnswers: Boolean = true): Unit = {
     SparkPlanTest.checkAnswer(
-        input, planFunction, expectedPlanFunction, sortAnswers, sqlContext) match {
+        input, planFunction, expectedPlanFunction, sortAnswers, _sqlContext) match {
       case Some(errorMessage) => fail(errorMessage)
       case None =>
     }
@@ -150,13 +157,13 @@ object SparkPlanTest {
       planFunction: SparkPlan => SparkPlan,
       expectedPlanFunction: SparkPlan => SparkPlan,
       sortAnswers: Boolean,
-      sqlContext: SQLContext): Option[String] = {
+      _sqlContext: SQLContext): Option[String] = {
 
     val outputPlan = planFunction(input.queryExecution.sparkPlan)
     val expectedOutputPlan = expectedPlanFunction(input.queryExecution.sparkPlan)
 
     val expectedAnswer: Seq[Row] = try {
-      executePlan(expectedOutputPlan, sqlContext)
+      executePlan(expectedOutputPlan, _sqlContext)
     } catch {
       case NonFatal(e) =>
         val errorMessage =
@@ -171,7 +178,7 @@ object SparkPlanTest {
     }
 
     val actualAnswer: Seq[Row] = try {
-      executePlan(outputPlan, sqlContext)
+      executePlan(outputPlan, _sqlContext)
     } catch {
       case NonFatal(e) =>
         val errorMessage =
@@ -211,12 +218,12 @@ object SparkPlanTest {
       planFunction: Seq[SparkPlan] => SparkPlan,
       expectedAnswer: Seq[Row],
       sortAnswers: Boolean,
-      sqlContext: SQLContext): Option[String] = {
+      _sqlContext: SQLContext): Option[String] = {
 
     val outputPlan = planFunction(input.map(_.queryExecution.sparkPlan))
 
     val sparkAnswer: Seq[Row] = try {
-      executePlan(outputPlan, sqlContext)
+      executePlan(outputPlan, _sqlContext)
     } catch {
       case NonFatal(e) =>
         val errorMessage =
@@ -279,10 +286,10 @@ object SparkPlanTest {
     }
   }
 
-  private def executePlan(outputPlan: SparkPlan, sqlContext: SQLContext): Seq[Row] = {
+  private def executePlan(outputPlan: SparkPlan, _sqlContext: SQLContext): Seq[Row] = {
     // A very simple resolver to make writing tests easier. In contrast to the real resolver
     // this is always case sensitive and does not try to handle scoping or complex type resolution.
-    val resolvedPlan = sqlContext.prepareForExecution.execute(
+    val resolvedPlan = _sqlContext.prepareForExecution.execute(
       outputPlan transform {
         case plan: SparkPlan =>
           val inputMap = plan.children.flatMap(_.output).map(a => (a.name, a)).toMap

sql/core/src/test/scala/org/apache/spark/sql/execution/debug/DebuggingSuite.scala

@@ -21,7 +21,7 @@ import org.apache.spark.SparkFunSuite
 import org.apache.spark.sql.test.MyTestSQLContext
 
 class DebuggingSuite extends SparkFunSuite with MyTestSQLContext {
-  private val ctx = sqlContextWithData
+  private val ctx = sqlContext
 
   test("DataFrame.debug()") {
     ctx.testData.debug()

sql/core/src/test/scala/org/apache/spark/sql/json/JsonSuite.scala

@@ -32,7 +32,7 @@ import org.apache.spark.sql.types._
 import org.apache.spark.util.Utils
 
 class JsonSuite extends QueryTest with TestJsonData with MyTestSQLContext {
-  private val _ctx = sqlContextWithData
+  private val _ctx = sqlContext
   import _ctx.implicits._
   import _ctx._
 

sql/core/src/test/scala/org/apache/spark/sql/parquet/ParquetCompatibilityTest.scala

@@ -16,18 +16,35 @@
  */
 
 package org.apache.spark.sql.parquet
+
 import java.io.File
 
 import scala.collection.JavaConversions._
 
+import org.scalatest.BeforeAndAfterAll
 import org.apache.hadoop.fs.Path
 import org.apache.parquet.hadoop.ParquetFileReader
 import org.apache.parquet.schema.MessageType
 
 import org.apache.spark.sql.QueryTest
 import org.apache.spark.util.Utils
 
-abstract class ParquetCompatibilityTest extends QueryTest with ParquetTest {
+/**
+ * Helper class for testing Parquet compatibility.
+ */
+private[sql] abstract class ParquetCompatibilityTest
+  extends AbstractParquetCompatibilityTest
+  with ParquetTest
+
+/**
+ * Abstract helper class for testing Parquet compatibility with a pluggable
+ * [[org.apache.spark.sql.SQLContext]].
+ */
+private[sql] abstract class AbstractParquetCompatibilityTest
+  extends QueryTest
+  with AbstractParquetTest
+  with BeforeAndAfterAll {
+
   protected var parquetStore: File = _
 
   /**

sql/core/src/test/scala/org/apache/spark/sql/parquet/ParquetTest.scala

@@ -23,8 +23,8 @@ import scala.reflect.ClassTag
 import scala.reflect.runtime.universe.TypeTag
 
 import org.apache.spark.SparkFunSuite
-import org.apache.spark.sql.test.SQLTestUtils
-import org.apache.spark.sql.{DataFrame, SaveMode}
+import org.apache.spark.sql.test.{AbstractSQLTestUtils, SQLTestUtils}
+import org.apache.spark.sql.{DataFrame, SaveMode, SQLContext}
 
 /**
  * A helper trait that provides convenient facilities for Parquet testing.
@@ -33,7 +33,13 @@ import org.apache.spark.sql.{DataFrame, SaveMode}
  * convenient to use tuples rather than special case classes when writing test cases/suites.
  * Especially, `Tuple1.apply` can be used to easily wrap a single type/value.
  */
-private[sql] trait ParquetTest extends SparkFunSuite with SQLTestUtils {
+private[sql] trait ParquetTest extends AbstractParquetTest with SQLTestUtils
+
+/**
+ * Abstract helper trait for Parquet tests with a pluggable [[SQLContext]].
+ */
+private[sql] trait AbstractParquetTest extends SparkFunSuite with AbstractSQLTestUtils {
+  protected def _sqlContext: SQLContext
 
   /**
    * Writes `data` to a Parquet file, which is then passed to `f` and will be deleted after `f`
@@ -43,7 +49,7 @@ private[sql] trait ParquetTest extends SparkFunSuite with SQLTestUtils {
       (data: Seq[T])
       (f: String => Unit): Unit = {
     withTempPath { file =>
-      sqlContext.createDataFrame(data).write.parquet(file.getCanonicalPath)
+      _sqlContext.createDataFrame(data).write.parquet(file.getCanonicalPath)
       f(file.getCanonicalPath)
     }
   }
@@ -55,7 +61,7 @@ private[sql] trait ParquetTest extends SparkFunSuite with SQLTestUtils {
   protected def withParquetDataFrame[T <: Product: ClassTag: TypeTag]
       (data: Seq[T])
       (f: DataFrame => Unit): Unit = {
-    withParquetFile(data)(path => f(sqlContext.read.parquet(path)))
+    withParquetFile(data)(path => f(_sqlContext.read.parquet(path)))
   }
 
   /**
@@ -67,14 +73,14 @@ private[sql] trait ParquetTest extends SparkFunSuite with SQLTestUtils {
       (data: Seq[T], tableName: String)
       (f: => Unit): Unit = {
     withParquetDataFrame(data) { df =>
-      sqlContext.registerDataFrameAsTable(df, tableName)
+      _sqlContext.registerDataFrameAsTable(df, tableName)
       withTempTable(tableName)(f)
     }
   }
 
   protected def makeParquetFile[T <: Product: ClassTag: TypeTag](
       data: Seq[T], path: File): Unit = {
-    sqlContext.createDataFrame(data).write.mode(SaveMode.Overwrite).parquet(path.getCanonicalPath)
+    _sqlContext.createDataFrame(data).write.mode(SaveMode.Overwrite).parquet(path.getCanonicalPath)
   }
 
   protected def makeParquetFile[T <: Product: ClassTag: TypeTag](

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

@@ -64,45 +64,36 @@ private[spark] class MyLocalSQLContext(sc: SparkContext) extends SQLContext(sc)
 private[spark] trait MyTestSQLContext extends SparkFunSuite with BeforeAndAfterAll {
 
   /**
-   * The [[SQLContext]] to use for all tests in this suite.
+   * The [[MyLocalSQLContext]] to use for all tests in this suite.
    *
    * By default, the underlying [[SparkContext]] will be run in local mode with the default
    * test configurations.
    */
-  private var _ctx: SQLContext = new MyLocalSQLContext
-
-  /** The [[SQLContext]] to use for all tests in this suite. */
-  protected def sqlContext: SQLContext = _ctx
+  private var _ctx: MyLocalSQLContext = new MyLocalSQLContext
 
   /**
    * The [[MyLocalSQLContext]] to use for all tests in this suite.
-   * This one comes with all the data prepared in advance.
    */
-  protected def sqlContextWithData: MyLocalSQLContext = {
-    _ctx match {
-      case local: MyLocalSQLContext => local
-      case _ => fail("this SQLContext does not have data prepared in advance")
-    }
-  }
+  protected def sqlContext: MyLocalSQLContext = _ctx
 
   /**
-   * Switch to the provided [[SQLContext]].
+   * Switch to the provided [[MyLocalSQLContext]].
    *
    * This stops the underlying [[SparkContext]] and expects a new one to be created.
    * This is needed because only one [[SparkContext]] is allowed per JVM.
    */
-  protected def switchSQLContext(newContext: () => SQLContext): Unit = {
+  protected def switchSQLContext(newContext: () => MyLocalSQLContext): Unit = {
     if (_ctx != null) {
       _ctx.sparkContext.stop()
       _ctx = newContext()
     }
   }
 
   /**
-   * Execute the given block of code with a custom [[SQLContext]].
-   * At the end of the method, a [[MyLocalSQLContext]] will be restored.
+   * Execute the given block of code with a custom [[MyLocalSQLContext]].
+   * At the end of the method, the default [[MyLocalSQLContext]] will be restored.
    */
-  protected def withSQLContext[T](newContext: () => SQLContext)(body: => T) {
+  protected def withSQLContext[T](newContext: () => MyLocalSQLContext)(body: => T) {
     switchSQLContext(newContext)
     try {
       body

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

@@ -26,6 +26,9 @@ import org.apache.spark.SparkFunSuite
 import org.apache.spark.sql.SQLContext
 import org.apache.spark.util.Utils
 
+/**
+ * General helper trait for common functionality in SQL tests.
+ */
 private[spark] trait SQLTestUtils
   extends SparkFunSuite
   with AbstractSQLTestUtils
@@ -34,6 +37,9 @@ private[spark] trait SQLTestUtils
   protected final override def _sqlContext = sqlContext
 }
 
+/**
+ * Abstract helper trait for SQL tests with a pluggable [[SQLContext]].
+ */
 private[spark] trait AbstractSQLTestUtils { this: SparkFunSuite =>
   protected def _sqlContext: SQLContext