diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
index 6dde79f74d3d8512b93cede7893458878b65b54f..39c0a2a0de0451544a2c463ae0035c6a9440f806 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
@@ -868,29 +868,27 @@ class TungstenAggregationQueryWithControlledFallbackSuite extends AggregationQue
 
   override protected def checkAnswer(actual: => DataFrame, expectedAnswer: Seq[Row]): Unit = {
     (0 to 2).foreach { fallbackStartsAt =>
-      sqlContext.setConf(
-        "spark.sql.TungstenAggregate.testFallbackStartsAt",
-        fallbackStartsAt.toString)
-
-      // Create a new df to make sure its physical operator picks up
-      // spark.sql.TungstenAggregate.testFallbackStartsAt.
-      // todo: remove it?
-      val newActual = DataFrame(sqlContext, actual.logicalPlan)
-
-      QueryTest.checkAnswer(newActual, expectedAnswer) match {
-        case Some(errorMessage) =>
-          val newErrorMessage =
-            s"""
-              |The following aggregation query failed when using TungstenAggregate with
-              |controlled fallback (it falls back to sort-based aggregation once it has processed
-              |$fallbackStartsAt input rows). The query is
-              |${actual.queryExecution}
-              |
-              |$errorMessage
-            """.stripMargin
-
-          fail(newErrorMessage)
-        case None =>
+      withSQLConf("spark.sql.TungstenAggregate.testFallbackStartsAt" -> fallbackStartsAt.toString) {
+        // Create a new df to make sure its physical operator picks up
+        // spark.sql.TungstenAggregate.testFallbackStartsAt.
+        // todo: remove it?
+        val newActual = DataFrame(sqlContext, actual.logicalPlan)
+
+        QueryTest.checkAnswer(newActual, expectedAnswer) match {
+          case Some(errorMessage) =>
+            val newErrorMessage =
+              s"""
+                |The following aggregation query failed when using TungstenAggregate with
+                |controlled fallback (it falls back to sort-based aggregation once it has processed
+                |$fallbackStartsAt input rows). The query is
+                |${actual.queryExecution}
+                |
+                |$errorMessage
+              """.stripMargin
+
+            fail(newErrorMessage)
+          case None =>
+        }
       }
     }
   }