Add missing SparkSessionTestWrapper for unsafe

mrpowers-io · Oct 5, 2024 · 9ea8d77 · 9ea8d77
1 parent dda5ef2
commit 9ea8d77
Showing 1 changed file with 21 additions and 0 deletions.
diff --git a/unsafe/src/test/scala/org/apache/spark/sql/utils/SparkSessionTestWrapper.scala b/unsafe/src/test/scala/org/apache/spark/sql/utils/SparkSessionTestWrapper.scala
@@ -0,0 +1,21 @@
+package org.apache.spark.sql.daria
+
+import org.apache.spark.sql.SparkSession
+
+trait SparkSessionTestWrapper {
+
+  lazy val spark: SparkSession = {
+    val session = SparkSession
+      .builder()
+      .master("local")
+      .appName("spark session")
+      .config(
+        "spark.sql.shuffle.partitions",
+        "1"
+      )
+      .getOrCreate()
+    session.sparkContext.setLogLevel("ERROR")
+    session
+  }
+
+}