From 9ea8d7723b1900b1ba5dfef87de762eb485c4e5f Mon Sep 17 00:00:00 2001
From: Tuan Pham <Tuan.Pham@wisetechglobal.com>
Date: Sat, 5 Oct 2024 22:46:36 +1000
Subject: [PATCH] Add missing SparkSessionTestWrapper for unsafe

---
 .../sql/utils/SparkSessionTestWrapper.scala   | 21 +++++++++++++++++++
 1 file changed, 21 insertions(+)
 create mode 100644 unsafe/src/test/scala/org/apache/spark/sql/utils/SparkSessionTestWrapper.scala

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
new file mode 100644
index 0000000..bc0706c
--- /dev/null
+++ 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
+  }
+
+}