A bit more

tests/integration/test_validate_dataframe_operator.py

@@ -12,7 +12,7 @@
 from integration.conftest import is_valid_gx_cloud_url, rand_name
 
 if TYPE_CHECKING:
-    import pyspark.sql as pyspark
+    from pyspark.sql import DataFrame as SparkDataFrame, SparkSession
     from pyspark.sql.connect.dataframe import DataFrame as SparkConnectDataFrame
     from pyspark.sql.connect.session import SparkSession as SparkConnectSession
 
@@ -60,7 +60,9 @@ def configure_dataframe() -> pd.DataFrame:
         assert is_valid_gx_cloud_url(result["result_url"])
 
     @pytest.mark.spark_integration
-    def test_spark(self, spark_session: pyspark.SparkSession) -> None:
+    def test_spark(self, spark_session: SparkSession) -> None:
+        import pyspark.sql as pyspark
+
         column_name = "col_A"
         task_id = f"test_spark_{rand_name()}"
 
@@ -115,7 +117,7 @@ def configure_dataframe() -> SparkConnectDataFrame:
 
 
 @pytest.fixture
-def spark_session() -> pyspark.SparkSession:
+def spark_session() -> SparkSession:
     import pyspark.sql as pyspark
 
     session = pyspark.SparkSession.builder.getOrCreate()