From e22db7acc5d20710010257fe463a3129eec10199 Mon Sep 17 00:00:00 2001 From: Tyler Hoffman Date: Tue, 14 Jan 2025 11:54:21 -0500 Subject: [PATCH] A bit more --- tests/integration/test_validate_dataframe_operator.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/integration/test_validate_dataframe_operator.py b/tests/integration/test_validate_dataframe_operator.py index 2f87e6f..2251d0f 100644 --- a/tests/integration/test_validate_dataframe_operator.py +++ b/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()