use sqlContext from LocalSparkContext instead of TestSQLContext

Author: mengxr

mllib/src/test/scala/org/apache/spark/ml/classification/LogisticRegressionSuite.scala

@@ -17,22 +17,17 @@
 
 package org.apache.spark.ml.classification
 
-import org.scalatest.{BeforeAndAfterAll, FunSuite}
+import org.scalatest.FunSuite
 
 import org.apache.spark.mllib.classification.LogisticRegressionSuite.generateLogisticInput
+import org.apache.spark.mllib.util.LocalSparkContext
 import org.apache.spark.sql.SchemaRDD
-import org.apache.spark.sql.test.TestSQLContext._
 
-class LogisticRegressionSuite extends FunSuite with BeforeAndAfterAll with Serializable {
+class LogisticRegressionSuite extends FunSuite with LocalSparkContext {
 
-  @transient var dataset: SchemaRDD = _
+  import sqlContext._
 
-  override def beforeAll(): Unit = {
-    super.beforeAll()
-    val points = generateLogisticInput(1.0, 1.0, 100, 42)
-    val rdd = sparkContext.parallelize(points)
-    dataset = createSchemaRDD(rdd)
-  }
+  val dataset: SchemaRDD = sc.parallelize(generateLogisticInput(1.0, 1.0, 100, 42), 2)
 
   test("logistic regression") {
     val lr = new LogisticRegression

mllib/src/test/scala/org/apache/spark/ml/tuning/CrossValidatorSuite.scala

@@ -17,22 +17,19 @@
 
 package org.apache.spark.ml.tuning
 
-import org.scalatest.{BeforeAndAfterAll, FunSuite}
+import org.scalatest.FunSuite
 
 import org.apache.spark.ml.classification.LogisticRegression
 import org.apache.spark.ml.evaluation.BinaryClassificationEvaluator
 import org.apache.spark.mllib.classification.LogisticRegressionSuite.generateLogisticInput
+import org.apache.spark.mllib.util.LocalSparkContext
 import org.apache.spark.sql.SchemaRDD
-import org.apache.spark.sql.test.TestSQLContext._
 
-class CrossValidatorSuite extends FunSuite with BeforeAndAfterAll with Serializable {
+class CrossValidatorSuite extends FunSuite with LocalSparkContext {
 
-  var dataset: SchemaRDD = _
+  import sqlContext._
 
-  override def beforeAll(): Unit = {
-    val points = generateLogisticInput(1.0, 1.0, 100, 42)
-    dataset = sparkContext.parallelize(points, 2)
-  }
+  val dataset: SchemaRDD = sc.parallelize(generateLogisticInput(1.0, 1.0, 100, 42), 2)
 
   test("cross validation with logistic regression") {
     val lr = new LogisticRegression

mllib/src/test/scala/org/apache/spark/mllib/util/LocalSparkContext.scala

@@ -17,26 +17,17 @@
 
 package org.apache.spark.mllib.util
 
-import org.scalatest.Suite
-import org.scalatest.BeforeAndAfterAll
+import org.scalatest.{BeforeAndAfterAll, Suite}
 
-import org.apache.spark.{SparkConf, SparkContext}
+import org.apache.spark.SparkContext
+import org.apache.spark.sql.SQLContext
 
 trait LocalSparkContext extends BeforeAndAfterAll { self: Suite =>
-  @transient var sc: SparkContext = _
-
-  override def beforeAll() {
-    val conf = new SparkConf()
-      .setMaster("local")
-      .setAppName("test")
-    sc = new SparkContext(conf)
-    super.beforeAll()
-  }
+  @transient val sc = new SparkContext("local[2]", "test")
+  @transient lazy val sqlContext = new SQLContext(sc)
 
   override def afterAll() {
-    if (sc != null) {
-      sc.stop()
-    }
+    sc.stop()
     super.afterAll()
   }
 }