Update workflows & pytest to choose the Sparksession instance based o…

…n a condition

.github/workflows/ci.yml

@@ -68,10 +68,11 @@ jobs:
       - name: Run tests with pytest against PySpark ${{ matrix.pyspark-version }}
         run: make test
 
-      - name: Run tests using Spark-Connect
+      - name: Run tests using Spark-Connect against PySpark ${{ matrix.pyspark-version }}
         env:
           HADOOP_VERSION: 3
           SPARK_VERSION: ${{ matrix.pyspark-version }}
+          SPARK_CONNECT_MODE_ENABLE: 1
         run: |
           if [[ "${SPARK_VERSION}" > "3.4" ]]; then
             sh scripts/run_spark_connect_server.sh

tests/spark.py

@@ -1,3 +1,7 @@
+import os
 from pyspark.sql import SparkSession
 
-spark = SparkSession.builder.master("local").appName("chispa").getOrCreate()
+if "SPARK_CONNECT_MODE_ENABLE" in os.environ:
+    spark = SparkSession.builder.remote("sc://localhost:15002").getOrCreate()
+else:
+    spark = SparkSession.builder.master("local").appName("chispa").getOrCreate()

tests/test_spark_connect.py

@@ -1,17 +1,11 @@
-import pytest
 import chispa
-from pyspark.sql import SparkSession
 from pyspark.sql.types import IntegerType, StringType, StructField, StructType
 
 import quinn
+from .spark import spark
 
 
-@pytest.fixture(scope="module")
-def spark():
-    return SparkSession.builder.remote("sc://localhost:15002").getOrCreate()
-
-
-def test_create_df(spark):
+def test_create_df():
     rows_data = [("abc", 1), ("lu", 2), ("torrence", 3)]
     col_specs = [("name", StringType()), ("age", IntegerType())]