From abb435b409b620ae71dd868a5647256eeb4194b4 Mon Sep 17 00:00:00 2001
From: Danny Meijer <10511979+dannymeijer@users.noreply.github.com>
Date: Tue, 12 Nov 2024 10:54:02 +0100
Subject: [PATCH] quick fix #2

---
 src/koheesio/spark/utils/common.py | 10 ++++------
 tests/spark/test_spark_utils.py    |  2 +-
 2 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/src/koheesio/spark/utils/common.py b/src/koheesio/spark/utils/common.py
index 3d0d79e..de918a7 100644
--- a/src/koheesio/spark/utils/common.py
+++ b/src/koheesio/spark/utils/common.py
@@ -79,18 +79,16 @@ def get_spark_minor_version() -> float:
 
 def check_if_pyspark_connect_is_supported() -> bool:
     """Check if the current version of PySpark supports the connect module"""
-    result = False
-
     if SPARK_MINOR_VERSION >= 3.5:
         try:
             importlib.import_module("pyspark.sql.connect")
             from pyspark.sql.connect.column import Column
 
-            _col: Column
-            result = True
+            _col: Column  # type: ignore
+            return True
         except (ModuleNotFoundError, ImportError):
-            result = False
-    return result
+            return False
+    return False
 
 
 if check_if_pyspark_connect_is_supported():
diff --git a/tests/spark/test_spark_utils.py b/tests/spark/test_spark_utils.py
index db76351..dfb4f9b 100644
--- a/tests/spark/test_spark_utils.py
+++ b/tests/spark/test_spark_utils.py
@@ -32,7 +32,7 @@ def test_unhappy_get_active_session_spark_connect(self):
             patch("pyspark.sql.SparkSession.getActiveSession", return_value=None),
         ):
             session = MagicMock(SparkSession=MagicMock(getActiveSession=MagicMock(return_value=None)))
-            with patch.dict("sys.modules", {"pyspark.sql.connect": session}):
+            with patch.dict("sys.modules", {"pyspark.sql.connect.session": session}):
                 with pytest.raises(
                     RuntimeError,
                     match="No active Spark session found. Please create a Spark session before using module "