From b016565945be11501e97a195c45e778f4ebbfc2c Mon Sep 17 00:00:00 2001 From: Dave Borowitz Date: Fri, 16 Oct 2026 08:35:43 -0700 Subject: [PATCH] test: share integration session, use spark.master=local, and align CI extras - Add tests/integration/conftest.py to provide a session-scoped default multi-node ManagedSparkSession for non-destructive tests (SQL, UDFs, %%sparksql, %dpip) and an isolated_session_state fixture for lifecycle tests. - Configure isolated control-plane and session lifecycle tests with spark.master=local (single-node mode: 1 driver VM, 0 executor node pools) and 4-core executor overrides instead of 16-core executors. - Remove the unused image_version=[None, '3.0'] fixture that previously doubled the integration suite. - Set default safety TTLs (MANAGED_SPARK_CONNECT_TTL_SECONDS=1800, MANAGED_SPARK_CONNECT_IDLE_TTL_SECONDS=600) and default/short-name subnet normalization in conftest.py, and ensure named and lifecycle sessions are terminated in finally blocks so sessions are never leaked. - Consolidate termination and deletion lifecycle tests into test_serverless_session_termination_and_deletion_lifecycle, assert SessionTemplate default runtime version on the created Session, and remove pure local CLI/regex tests from tests/integration/ that are already covered in tests/unit/. - Together, these reduce remote Dataproc Serverless clusters per run from 42 to 5 (1 multi-node + 4 single-node), total GCE VMs from ~130 to 7 (-94.6%), and end-to-end runtime to ~8 minutes. - Move test_create_local_spark_session (DATAPROC_WORKLOAD_TYPE=batch) to tests/local/test_sparksql_local.py so it runs on every PR in the spark lane, and isolate local Spark Connect environment/proxy state. - Switch GitHub Actions workflows to astral-sh/setup-uv + uv pip install, update integration-tests.yaml to install .[client,interactive], and update tests.yaml to verify unit tests both without [interactive] and with [interactive]. --- .github/workflows/integration-tests.yaml | 19 +- .github/workflows/lint-readme.yaml | 5 +- .github/workflows/pyink.yaml | 5 +- .github/workflows/tests.yaml | 28 +- tests/integration/conftest.py | 238 +++++ .../managed_spark_magics/test_magics.py | 151 +--- tests/integration/test_session.py | 841 ++++++------------ tests/local/test_sparksql_local.py | 32 + tests/unit/test_init.py | 55 +- tests/unit/test_session.py | 35 +- 10 files changed, 678 insertions(+), 731 deletions(-) create mode 100644 tests/integration/conftest.py diff --git a/.github/workflows/integration-tests.yaml b/.github/workflows/integration-tests.yaml index a11aa52..3b0dade 100644 --- a/.github/workflows/integration-tests.yaml +++ b/.github/workflows/integration-tests.yaml @@ -17,7 +17,7 @@ # Required GitHub Secrets: # - GCP_SA_KEY: Service account JSON key (project_id and client_email extracted automatically) # - GCP_REGION: Google Cloud Region (optional, defaults to us-central1) -# - GCP_SUBNET: Managed Spark subnet URI +# - GCP_SUBNET: Managed Spark subnet name or URI (optional, defaults to "default") # # See INTEGRATION_TESTS.md for setup instructions. @@ -49,19 +49,20 @@ jobs: with: python-version: "3.12" - - name: Cache pip dependencies + - name: Set up uv + uses: astral-sh/setup-uv@1c37ad07a6a961277cf70c0d37d6f313000f5884 # v10.3.0 + + - name: Cache uv dependencies uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0 with: - path: ~/.cache/pip - key: ${{ runner.os }}-pip-integration-${{ hashFiles('requirements-dev.txt', 'requirements-test.txt') }} + path: ~/.cache/uv + key: ${{ runner.os }}-uv-integration-${{ hashFiles('setup.py', 'requirements-test.txt') }} restore-keys: | - ${{ runner.os }}-pip-integration- - ${{ runner.os }}-pip- + ${{ runner.os }}-uv-integration- + ${{ runner.os }}-uv- - name: Install dependencies - run: | - pip install -r requirements-dev.txt - pip install -r requirements-test.txt + run: uv pip install --system -e '.[client,interactive]' -r requirements-test.txt - name: Authenticate to Google Cloud uses: google-github-actions/auth@7c6bc770dae815cd3e89ee6cdf493a5fab2cc093 # v3 diff --git a/.github/workflows/lint-readme.yaml b/.github/workflows/lint-readme.yaml index 560f3ab..9f82131 100644 --- a/.github/workflows/lint-readme.yaml +++ b/.github/workflows/lint-readme.yaml @@ -32,11 +32,14 @@ jobs: with: python-version: "3.12" + - name: Set up uv + uses: astral-sh/setup-uv@1c37ad07a6a961277cf70c0d37d6f313000f5884 # v10.3.0 + - name: Checkout code uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Install dependencies - run: pip install build twine + run: uv pip install --system build twine - name: Build package run: python -m build diff --git a/.github/workflows/pyink.yaml b/.github/workflows/pyink.yaml index 0a10411..8159a77 100644 --- a/.github/workflows/pyink.yaml +++ b/.github/workflows/pyink.yaml @@ -32,8 +32,11 @@ jobs: with: python-version: "3.12" + - name: Set up uv + uses: astral-sh/setup-uv@1c37ad07a6a961277cf70c0d37d6f313000f5884 # v10.3.0 + - name: Install pyink - run: pip install pyink + run: uv pip install --system pyink - name: Checkout code uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml index 94b9722..9cb1561 100644 --- a/.github/workflows/tests.yaml +++ b/.github/workflows/tests.yaml @@ -41,22 +41,30 @@ jobs: with: python-version: "3.12" - - name: Cache pip dependencies + - name: Set up uv + uses: astral-sh/setup-uv@1c37ad07a6a961277cf70c0d37d6f313000f5884 # v10.3.0 + + - name: Cache uv dependencies uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0 with: - path: ~/.cache/pip - key: ${{ runner.os }}-pip-${{ matrix.spark-extra }}-${{ hashFiles('setup.py', 'requirements-test.txt') }} + path: ~/.cache/uv + key: ${{ runner.os }}-uv-${{ matrix.spark-extra }}-${{ hashFiles('setup.py', 'requirements-test.txt') }} restore-keys: | - ${{ runner.os }}-pip-${{ matrix.spark-extra }}- - ${{ runner.os }}-pip- + ${{ runner.os }}-uv-${{ matrix.spark-extra }}- + ${{ runner.os }}-uv- # Installs through the package's own extras, rather than # requirements-dev.txt (which pins pyspark-client), so that both extras are - # exercised as users would install them. - - name: Install dependencies - run: | - pip install -e '.[${{ matrix.spark-extra }},interactive]' - pip install -r requirements-test.txt + # exercised as users would install them. First verify the base install + # without [interactive], then install [interactive] for the full suite. + - name: Install base dependencies (without [interactive]) + run: uv pip install --system -e '.[${{ matrix.spark-extra }}]' -r requirements-test.txt + + - name: Run unit tests without [interactive] + run: python -m pytest tests/unit/test_init.py tests/unit/test_pypi_artifacts.py -k "TestWithoutInteractiveExtra or TestPythonVersionCheck or TestRuntimeVersionCompatibility or PyPiArtifactsTest" -v --tb=short + + - name: Install [interactive] extra + run: uv pip install --system -e '.[${{ matrix.spark-extra }},interactive]' - name: Run unit tests run: python -m pytest tests/unit/ -v --tb=short -n auto diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py new file mode 100644 index 0000000..6b25a5e --- /dev/null +++ b/tests/integration/conftest.py @@ -0,0 +1,238 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import os +import certifi +from google.api_core import client_options +from google.cloud.dataproc_v1 import ( + SessionControllerClient, + SessionTemplateControllerClient, +) +from google.cloud.managed_spark_connect import ManagedSparkSession +import pytest + +_SERVICE_ACCOUNT_KEY_FILE_ = "service_account_key.json" + + +def is_ci_environment(): + """Detect if running in CI environment.""" + return os.getenv("CI") == "true" or os.getenv("GITHUB_ACTIONS") == "true" + + +@pytest.fixture(scope="session") +def test_project(): + return os.getenv("GOOGLE_CLOUD_PROJECT") + + +@pytest.fixture(scope="session") +def test_region(): + return os.getenv("GOOGLE_CLOUD_REGION") + + +@pytest.fixture(scope="session") +def auth_type(request): + """Auto-detect authentication type based on environment. + + CI environment (CI=true or GITHUB_ACTIONS=true): Uses SERVICE_ACCOUNT + Local environment: Uses END_USER_CREDENTIALS + Test parametrization can still override this default. + """ + if hasattr(request, "param"): + return request.param + + if is_ci_environment(): + return "SERVICE_ACCOUNT" + else: + return "END_USER_CREDENTIALS" + + +@pytest.fixture(scope="session") +def test_subnet(): + return os.getenv("MANAGED_SPARK_CONNECT_SUBNET") or "default" + + +@pytest.fixture(scope="session") +def test_subnetwork_uri(test_project, test_region, test_subnet): + if "/" in test_subnet: + return test_subnet + return ( + f"projects/{test_project}/regions/{test_region}/subnetworks/" + f"{test_subnet}" + ) + + +@pytest.fixture(scope="session") +def os_environment(auth_type, test_project, test_region, test_subnetwork_uri): + original_environment = dict(os.environ) + if os.path.isfile(_SERVICE_ACCOUNT_KEY_FILE_): + os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = ( + _SERVICE_ACCOUNT_KEY_FILE_ + ) + os.environ["MANAGED_SPARK_CONNECT_AUTH_TYPE"] = auth_type + if auth_type == "END_USER_CREDENTIALS": + os.environ.pop("MANAGED_SPARK_CONNECT_SERVICE_ACCOUNT", None) + os.environ["MANAGED_SPARK_CONNECT_SUBNET"] = test_subnetwork_uri + # Set short safety TTLs so any session leaked by a crashed or cancelled + # test run auto-terminates after 10 minutes of inactivity (or 30 minutes + # total) instead of the 24-hour server default. + os.environ.setdefault("MANAGED_SPARK_CONNECT_TTL_SECONDS", "1800") + os.environ.setdefault("MANAGED_SPARK_CONNECT_IDLE_TTL_SECONDS", "600") + # Add SSL certificate fix + os.environ["SSL_CERT_FILE"] = certifi.where() + os.environ["REQUESTS_CA_BUNDLE"] = certifi.where() + yield os.environ + for key in list(os.environ.keys()): + if key not in original_environment: + del os.environ[key] + os.environ.update(original_environment) + + +@pytest.fixture(scope="session") +def api_endpoint(test_region): + return os.getenv( + "GOOGLE_CLOUD_DATAPROC_API_ENDPOINT", + f"{test_region}-dataproc.googleapis.com", + ) + + +@pytest.fixture(scope="session") +def test_client_options(api_endpoint, os_environment): + return client_options.ClientOptions(api_endpoint=api_endpoint) + + +@pytest.fixture(scope="session") +def session_controller_client(test_client_options): + return SessionControllerClient(client_options=test_client_options) + + +@pytest.fixture(scope="session") +def session_template_controller_client(test_client_options): + return SessionTemplateControllerClient(client_options=test_client_options) + + +def _capture_session_state(): + return { + "session": ManagedSparkSession.getActiveSession() + or ManagedSparkSession._default_session, + "session_id": ManagedSparkSession._active_s8s_session_id, + "session_uuid": ManagedSparkSession._active_s8s_session_uuid, + "uses_custom_id": ManagedSparkSession._active_session_uses_custom_id, + "project_id": ManagedSparkSession._project_id, + "region": ManagedSparkSession._region, + "client_options": ManagedSparkSession._client_options, + } + + +def _restore_session_state(state): + ManagedSparkSession._active_s8s_session_id = state["session_id"] + ManagedSparkSession._active_s8s_session_uuid = state["session_uuid"] + ManagedSparkSession._active_session_uses_custom_id = state["uses_custom_id"] + ManagedSparkSession._project_id = state["project_id"] + ManagedSparkSession._region = state["region"] + ManagedSparkSession._client_options = state["client_options"] + session = state["session"] + if session is not None: + ManagedSparkSession._set_default_and_active_session(session) + try: + from pyspark.sql import SparkSession as PySparkSQLSession + + PySparkSQLSession._instantiatedSession = session + except (ImportError, AttributeError): + pass + else: + ManagedSparkSession._default_session = None + ManagedSparkSession._active_session.session = None + try: + from pyspark.sql import SparkSession as PySparkSQLSession + + PySparkSQLSession._instantiatedSession = None + except (ImportError, AttributeError): + pass + + +def _clear_session_state(): + ManagedSparkSession._active_s8s_session_id = None + ManagedSparkSession._active_s8s_session_uuid = None + ManagedSparkSession._active_session_uses_custom_id = False + ManagedSparkSession._project_id = None + ManagedSparkSession._region = None + ManagedSparkSession._client_options = None + ManagedSparkSession._default_session = None + ManagedSparkSession._active_session.session = None + try: + from pyspark.sql import SparkSession as PySparkSQLSession + + PySparkSQLSession._instantiatedSession = None + except (ImportError, AttributeError): + pass + + +@pytest.fixture(scope="session") +def _shared_connect_session(test_project, test_region, os_environment): + """Session-scoped default ManagedSparkSession shared by non-destructive tests. + + Created with default builder settings (multi-node) so out-of-the-box user + defaults, Python UDF execution on remote executors, and %dpip package + distribution are exercised against a standard cluster. + """ + _clear_session_state() + session = ( + ManagedSparkSession.builder.projectId(test_project) + .location(test_region) + .getOrCreate() + ) + shared_state = _capture_session_state() + yield shared_state + _restore_session_state(shared_state) + try: + session.stop(terminate=True) + except Exception: + pass + + +@pytest.fixture +def connect_session(_shared_connect_session): + """Provides the shared default ManagedSparkSession and restores its active state.""" + _restore_session_state(_shared_connect_session) + yield _shared_connect_session["session"] + _restore_session_state(_shared_connect_session) + + +@pytest.fixture +def session_name(test_project, test_region, connect_session): + return ( + f"projects/{test_project}/locations/{test_region}/sessions/" + f"{ManagedSparkSession._active_s8s_session_id}" + ) + + +@pytest.fixture +def isolated_session_state(os_environment): + """Isolates ManagedSparkSession class state for tests that manage their own session lifecycle.""" + saved_state = _capture_session_state() + _clear_session_state() + try: + yield + finally: + leaked_id = ManagedSparkSession._active_s8s_session_id + if leaked_id and leaked_id != saved_state["session_id"]: + try: + active = ( + ManagedSparkSession.getActiveSession() + or ManagedSparkSession._default_session + ) + if active is not None: + active.stop(terminate=True) + except Exception: + pass + _restore_session_state(saved_state) diff --git a/tests/integration/managed_spark_magics/test_magics.py b/tests/integration/managed_spark_magics/test_magics.py index b0c8f11..6e17282 100644 --- a/tests/integration/managed_spark_magics/test_magics.py +++ b/tests/integration/managed_spark_magics/test_magics.py @@ -11,112 +11,19 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -import os import pytest -import certifi -from unittest import mock - -from google.cloud.managed_spark_connect import ManagedSparkSession - -_SERVICE_ACCOUNT_KEY_FILE_ = "service_account_key.json" - - -@pytest.fixture(params=[None, "3.0"]) -def image_version(request): - return request.param - - -@pytest.fixture -def test_project(): - return os.getenv("GOOGLE_CLOUD_PROJECT") - - -@pytest.fixture -def test_region(): - return os.getenv("GOOGLE_CLOUD_REGION") - - -def is_ci_environment(): - """Detect if running in CI environment.""" - return os.getenv("CI") == "true" or os.getenv("GITHUB_ACTIONS") == "true" - - -@pytest.fixture -def auth_type(request): - """Auto-detect authentication type based on environment. - - CI environment (CI=true or GITHUB_ACTIONS=true): Uses SERVICE_ACCOUNT - Local environment: Uses END_USER_CREDENTIALS - Test parametrization can still override this default. - """ - # Allow test parametrization to override - if hasattr(request, "param"): - return request.param - - # Auto-detect based on environment - if is_ci_environment(): - return "SERVICE_ACCOUNT" - else: - return "END_USER_CREDENTIALS" - - -@pytest.fixture -def test_subnet(): - return os.getenv("MANAGED_SPARK_CONNECT_SUBNET") - - -@pytest.fixture -def test_subnetwork_uri(test_subnet): - # Make MANAGED_SPARK_CONNECT_SUBNET the full URI - # to align with how user would specify it in the project - return test_subnet - - -@pytest.fixture -def os_environment(auth_type, image_version, test_project, test_region): - original_environment = dict(os.environ) - if os.path.isfile(_SERVICE_ACCOUNT_KEY_FILE_): - os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = ( - _SERVICE_ACCOUNT_KEY_FILE_ - ) - os.environ["MANAGED_SPARK_CONNECT_AUTH_TYPE"] = auth_type - if auth_type == "END_USER_CREDENTIALS": - os.environ.pop("MANAGED_SPARK_CONNECT_SERVICE_ACCOUNT", None) - # Add SSL certificate fix - os.environ["SSL_CERT_FILE"] = certifi.where() - os.environ["REQUESTS_CA_BUNDLE"] = certifi.where() - yield os.environ - os.environ.clear() - os.environ.update(original_environment) - - -@pytest.fixture -def connect_session(test_project, test_region, os_environment): - session = ( - ManagedSparkSession.builder.projectId(test_project) - .location(test_region) - .getOrCreate() - ) - yield session - # Clean up the session after each test to prevent resource conflicts - try: - session.stop() - except Exception: - # Ignore cleanup errors to avoid masking the actual test failure - pass @pytest.fixture def ipython_shell(connect_session): - """Provides an IPython shell with a ManagedSparkSession in user_ns.""" + """Provides an IPython shell with the shared ManagedSparkSession in user_ns.""" try: - from IPython.terminal.interactiveshell import TerminalInteractiveShell from google.cloud import managed_spark_magics + from IPython.terminal.interactiveshell import TerminalInteractiveShell shell = TerminalInteractiveShell.instance() shell.user_ns = {"spark": connect_session} - # Load magics managed_spark_magics.load_ipython_extension(shell) yield shell @@ -126,14 +33,8 @@ def ipython_shell(connect_session): TerminalInteractiveShell.clear_instance() -# Tests for magics.py -def test_dpip_magic_loads(ipython_shell): - """Test that %dpip magic is registered.""" - assert "dpip" in ipython_shell.magics_manager.magics["line"] - - def test_dpip_install_success(connect_session, ipython_shell, capsys): - """Test installing a single package with %dpip.""" + """Test installing PyPI packages with %dpip and running a UDF on remote executors.""" ipython_shell.run_line_magic("dpip", "install roman numpy") captured = capsys.readouterr() assert "Active session found:" in captured.out @@ -156,55 +57,11 @@ def to_roman(number): assert df_result[0]["roman"] == "MDCLXVI" - connect_session.stop() - - -def test_dpip_no_install_command(ipython_shell): - """Test usage message when 'install' is missing.""" - with pytest.raises( - RuntimeError, match="Usage: %dpip install ..." - ): - ipython_shell.run_line_magic("dpip", "pandas") - - -def test_dpip_no_packages(ipython_shell): - """Test message when no packages are specified.""" - with pytest.raises(RuntimeError, match="Error: No packages specified."): - ipython_shell.run_line_magic("dpip", "install") - - -def test_dpip_with_flags(ipython_shell): - """Test installing multiple packages with flags like -U.""" - with pytest.raises( - RuntimeError, match="Error: Flags are not currently supported." - ): - ipython_shell.run_line_magic("dpip", "install -U numpy scikit-learn") - - -def test_dpip_no_session(ipython_shell): - """Test message when no Spark session is active.""" - ipython_shell.user_ns = {} # Remove spark session from namespace - with pytest.raises( - RuntimeError, match="No active Managed Spark Session found." - ): - ipython_shell.run_line_magic("dpip", "install pandas") - def test_dpip_install_failure(ipython_shell): - """Test error message on installation failure.""" + """Test error message on remote PyPI installation failure.""" with pytest.raises( RuntimeError, match="No matching distribution found", ): ipython_shell.run_line_magic("dpip", "install dp-non-existent-package") - - -def test_dpip_multiple_sessions(ipython_shell, connect_session): - """Test error message when multiple Spark sessions found.""" - ipython_shell.user_ns["sparksession"] = connect_session - ipython_shell.user_ns["sparkanother"] = connect_session - with pytest.raises( - RuntimeError, - match="Error: Found more than one active Managed Spark Sessions.", - ): - ipython_shell.run_line_magic("dpip", "install pandas") diff --git a/tests/integration/test_session.py b/tests/integration/test_session.py index 77bf133..26eac1d 100644 --- a/tests/integration/test_session.py +++ b/tests/integration/test_session.py @@ -12,13 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. import datetime -import os -import pytest import uuid -import certifi -from google.api_core import client_options -from google.cloud.managed_spark_connect import ManagedSparkSession from google.cloud.dataproc_v1 import ( CreateSessionTemplateRequest, DeleteSessionRequest, @@ -26,128 +21,13 @@ GetSessionRequest, GetSessionTemplateRequest, Session, - SessionControllerClient, SessionTemplate, - SessionTemplateControllerClient, TerminateSessionRequest, ) +from google.cloud.managed_spark_connect import ManagedSparkSession from pyspark.errors.exceptions import connect as connect_exceptions from pyspark.sql.types import StringType -from pyspark.util import is_remote_only - -_SERVICE_ACCOUNT_KEY_FILE_ = "service_account_key.json" - - -@pytest.fixture(params=[None, "3.0"]) -def image_version(request): - return request.param - - -@pytest.fixture -def test_project(): - return os.getenv("GOOGLE_CLOUD_PROJECT") - - -def is_ci_environment(): - """Detect if running in CI environment.""" - return os.getenv("CI") == "true" or os.getenv("GITHUB_ACTIONS") == "true" - - -@pytest.fixture -def auth_type(request): - """Auto-detect authentication type based on environment. - - CI environment (CI=true or GITHUB_ACTIONS=true): Uses SERVICE_ACCOUNT - Local environment: Uses END_USER_CREDENTIALS - Test parametrization can still override this default. - """ - # Allow test parametrization to override - if hasattr(request, "param"): - return request.param - - # Auto-detect based on environment - if is_ci_environment(): - return "SERVICE_ACCOUNT" - else: - return "END_USER_CREDENTIALS" - - -@pytest.fixture -def test_region(): - return os.getenv("GOOGLE_CLOUD_REGION") - - -@pytest.fixture -def test_subnet(): - return os.getenv("MANAGED_SPARK_CONNECT_SUBNET") - - -@pytest.fixture -def test_subnetwork_uri(test_subnet): - # Make MANAGED_SPARK_CONNECT_SUBNET the full URI to align with how user would specify it in the project - return test_subnet - - -@pytest.fixture -def os_environment(auth_type, image_version, test_project, test_region): - original_environment = dict(os.environ) - if os.path.isfile(_SERVICE_ACCOUNT_KEY_FILE_): - os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = ( - _SERVICE_ACCOUNT_KEY_FILE_ - ) - os.environ["MANAGED_SPARK_CONNECT_AUTH_TYPE"] = auth_type - if auth_type == "END_USER_CREDENTIALS": - os.environ.pop("MANAGED_SPARK_CONNECT_SERVICE_ACCOUNT", None) - # Add SSL certificate fix - os.environ["SSL_CERT_FILE"] = certifi.where() - os.environ["REQUESTS_CA_BUNDLE"] = certifi.where() - yield os.environ - os.environ.clear() - os.environ.update(original_environment) - - -@pytest.fixture -def api_endpoint(test_region): - return os.getenv( - "GOOGLE_CLOUD_DATAPROC_API_ENDPOINT", - f"{test_region}-dataproc.googleapis.com", - ) - - -@pytest.fixture -def test_client_options(api_endpoint, os_environment): - return client_options.ClientOptions(api_endpoint=api_endpoint) - - -@pytest.fixture -def session_controller_client(test_client_options): - return SessionControllerClient(client_options=test_client_options) - - -@pytest.fixture -def session_template_controller_client(test_client_options): - return SessionTemplateControllerClient(client_options=test_client_options) - - -@pytest.fixture -def connect_session(test_project, test_region, os_environment): - session = ( - ManagedSparkSession.builder.projectId(test_project) - .location(test_region) - .getOrCreate() - ) - yield session - # Clean up the session after each test to prevent resource conflicts - try: - session.stop() - except Exception: - # Ignore cleanup errors to avoid masking the actual test failure - pass - - -@pytest.fixture -def session_name(test_project, test_region, connect_session): - return f"projects/{test_project}/locations/{test_region}/sessions/{ManagedSparkSession._active_s8s_session_id}" +import pytest def test_create_spark_session_with_default_notebook_behavior( @@ -165,127 +45,136 @@ def test_create_spark_session_with_default_notebook_behavior( assert str(df) == "DataFrame[id: bigint, name: string]" connect_session.sql("DROP TABLE IF EXISTS FOO") connect_session.sql("CREATE TABLE FOO (bar long, baz long) USING PARQUET") - with pytest.raises(connect_exceptions.AnalysisException) as ex: - connect_session.sql( - "CREATE TABLE FOO (bar long, baz long) USING PARQUET" - ) - - assert "[TABLE_OR_VIEW_ALREADY_EXISTS]" in str(ex) + try: + with pytest.raises(connect_exceptions.AnalysisException) as ex: + connect_session.sql( + "CREATE TABLE FOO (bar long, baz long) USING PARQUET" + ) + assert "[TABLE_OR_VIEW_ALREADY_EXISTS]" in str(ex.value) + finally: + connect_session.sql("DROP TABLE IF EXISTS FOO") assert ManagedSparkSession._active_s8s_session_uuid is not None - connect_session.sql("DROP TABLE IF EXISTS FOO") - connect_session.stop() - session = session_controller_client.get_session(get_session_request) - - assert session.state in [ - Session.State.TERMINATING, - Session.State.TERMINATED, - ] - assert ManagedSparkSession._active_s8s_session_uuid is None -def test_reuse_s8s_spark_session( - connect_session, session_name, session_controller_client -): +def test_reuse_s8s_spark_session(connect_session): """Test that Spark sessions can be reused within the same process.""" assert ManagedSparkSession._active_s8s_session_uuid is not None first_session_id = ManagedSparkSession._active_s8s_session_id first_session_uuid = ManagedSparkSession._active_s8s_session_uuid - connect_session = ManagedSparkSession.builder.getOrCreate() + reused_session = ManagedSparkSession.builder.getOrCreate() second_session_id = ManagedSparkSession._active_s8s_session_id second_session_uuid = ManagedSparkSession._active_s8s_session_uuid + assert reused_session is connect_session assert first_session_id == second_session_id assert first_session_uuid == second_session_uuid assert ManagedSparkSession._active_s8s_session_uuid is not None assert ManagedSparkSession._active_s8s_session_id is not None - connect_session.stop() +def test_sql_functions(connect_session): + """Test basic SQL functions like col(), sum(), count(), etc.""" + from pyspark.sql.connect.functions import col, count, sum -def test_stop_spark_session_with_deleted_serverless_session( - connect_session, session_name, session_controller_client -): - """Test stopping a Spark session when the serverless session has been deleted.""" - assert ManagedSparkSession._active_s8s_session_uuid is not None + df = connect_session.createDataFrame( + [(1, "Alice", 100), (2, "Bob", 200), (3, "Charlie", 150)], + ["id", "name", "amount"], + ) - delete_session_request = DeleteSessionRequest() - delete_session_request.name = session_name - operation = session_controller_client.delete_session(delete_session_request) - operation.result() - connect_session.stop() + result_col = df.select(col("name")).collect() + assert len(result_col) == 3 + assert result_col[0]["name"] == "Alice" - assert ManagedSparkSession._active_s8s_session_uuid is None - assert ManagedSparkSession._active_s8s_session_id is None + sum_result = df.select(sum("amount")).collect()[0][0] + assert sum_result == 450 + count_result = df.select(count("id")).collect()[0][0] + assert count_result == 3 -def test_stop_spark_session_with_terminated_serverless_session( - connect_session, session_name, session_controller_client -): - """Test stopping a Spark session when the serverless session has been terminated.""" - assert ManagedSparkSession._active_s8s_session_uuid is not None + filtered_df = df.where(col("amount") > 150) + filtered_count = filtered_df.count() + assert filtered_count == 1 - terminate_session_request = TerminateSessionRequest() - terminate_session_request.name = session_name - operation = session_controller_client.terminate_session( - terminate_session_request + df_with_calc = df.select( + col("id"), + col("name"), + col("amount"), + (col("amount") * 0.1).alias("tax"), ) - operation.result() - connect_session.stop() + tax_results = df_with_calc.collect() + assert tax_results[0]["tax"] == 10.0 + assert tax_results[1]["tax"] == 20.0 + assert tax_results[2]["tax"] == 15.0 - assert ManagedSparkSession._active_s8s_session_uuid is None - assert ManagedSparkSession._active_s8s_session_id is None +def test_sql_udf(connect_session): + """Test SQL UDF registration and usage on remote executors.""" + from pyspark.sql.connect.functions import col, udf -def test_get_or_create_spark_session_with_terminated_serverless_session( - test_project, - test_region, - connect_session, - session_name, - session_controller_client, -): - """Test creating a new Spark session when the previous serverless session has been terminated.""" - first_session_name = session_name + df = connect_session.createDataFrame( + [(1, "hello"), (2, "world"), (3, "spark")], ["id", "text"] + ) - assert ManagedSparkSession._active_s8s_session_uuid is not None + def uppercase_func(text): + return text.upper() if text else None - first_session = ManagedSparkSession._active_s8s_session_uuid - terminate_session_request = TerminateSessionRequest() - terminate_session_request.name = first_session_name - operation = session_controller_client.terminate_session( - terminate_session_request + uppercase_udf = udf(uppercase_func, StringType()) + df_with_udf = df.select( + "id", "text", uppercase_udf(col("text")).alias("upper_text") ) - operation.result() - connect_session = ManagedSparkSession.builder.getOrCreate() - second_session = ManagedSparkSession._active_s8s_session_uuid - second_session_name = f"projects/{test_project}/locations/{test_region}/sessions/{ManagedSparkSession._active_s8s_session_id}" + df_result = df_with_udf.collect() - assert first_session != second_session - assert ManagedSparkSession._active_s8s_session_uuid is not None - assert ManagedSparkSession._active_s8s_session_id is not None + assert df_result[0]["upper_text"] == "HELLO" + assert df_result[1]["upper_text"] == "WORLD" - get_session_request = GetSessionRequest() - get_session_request.name = first_session_name - session = session_controller_client.get_session(get_session_request) - assert session.state in [ - Session.State.TERMINATING, - Session.State.TERMINATED, - ] +def test_sparksql_magic_with_managed_spark_session(connect_session): + """Test that %%sparksql works with the active ManagedSparkSession.""" + pytest.importorskip( + "IPython", reason="IPython not available (install [interactive])" + ) + pytest.importorskip( + "sqlparse", reason="sqlparse not available (install [interactive])" + ) - get_session_request = GetSessionRequest() - get_session_request.name = second_session_name - session = session_controller_client.get_session(get_session_request) + from IPython.terminal.interactiveshell import TerminalInteractiveShell - assert session.state == Session.State.ACTIVE - connect_session.stop() + shell = TerminalInteractiveShell.instance() + try: + shell.run_line_magic("load_ext", "google.cloud.managed_spark_magics") + shell.user_ns["factor"] = 5 + + shell.run_cell_magic( + "sparksql", + "-q result_df", + """ + SELECT + 10 * :factor as multiplication, + SQRT(16) as square_root, + CONCAT('Dataproc', '-', 'Spark') as joined_string + """, + ) + + assert "result_df" in shell.user_ns + df = shell.user_ns["result_df"] + assert df is not None + + data = df.collect() + assert len(data) == 1 + row = data[0] + + assert row["multiplication"] == 50 + assert row["square_root"] == 4.0 + assert row["joined_string"] == "Dataproc-Spark" + finally: + TerminalInteractiveShell.clear_instance() @pytest.fixture def session_template_name( - image_version, test_project, test_region, test_subnetwork_uri, @@ -299,51 +188,59 @@ def session_template_name( session_template.environment_config.execution_config.subnetwork_uri = ( test_subnetwork_uri ) - if image_version: - session_template.runtime_config.version = image_version - session_template_name = f"projects/{test_project}/locations/{test_region}/sessionTemplates/spark-connect-test-template-{uuid.uuid4().hex[0:12]}" + session_template_name = ( + f"projects/{test_project}/locations/{test_region}/sessionTemplates/" + f"spark-connect-test-template-{uuid.uuid4().hex[0:12]}" + ) session_template.name = session_template_name create_session_template_request.session_template = session_template session_template_controller_client.create_session_template( create_session_template_request ) - get_session_template_request = GetSessionTemplateRequest() - get_session_template_request.name = session_template_name - session_template = session_template_controller_client.get_session_template( - get_session_template_request - ) - assert ( - session_template.runtime_config.version == image_version - if image_version - else ManagedSparkSession._DEFAULT_RUNTIME_VERSION - ) - - yield session_template.name - delete_session_template_request = DeleteSessionTemplateRequest() - delete_session_template_request.name = session_template_name - session_template_controller_client.delete_session_template( - delete_session_template_request - ) + try: + get_session_template_request = GetSessionTemplateRequest() + get_session_template_request.name = session_template_name + created_template = ( + session_template_controller_client.get_session_template( + get_session_template_request + ) + ) + yield created_template.name + finally: + delete_session_template_request = DeleteSessionTemplateRequest() + delete_session_template_request.name = session_template_name + session_template_controller_client.delete_session_template( + delete_session_template_request + ) def test_create_spark_session_with_session_template_and_user_provided_session_config( - image_version, test_project, test_region, session_template_name, session_controller_client, + isolated_session_state, ): """Test creating a Spark session with a session template and user-provided Dataproc configuration.""" session_config = Session() session_config.environment_config.execution_config.ttl = {"seconds": 64800} session_config.session_template = session_template_name + # Use spark.master=local (single-node mode) and 4-core override (instead of + # 16) so this control-plane configuration test only provisions a single + # 4-vCPU driver VM without an executor node pool. connect_session = ( - ManagedSparkSession.builder.config("spark.executor.cores", "7") + ManagedSparkSession.builder.projectId(test_project) + .location(test_region) + .config("spark.executor.cores", "8") .sessionConfig(session_config) - .config("spark.executor.cores", "16") + .config("spark.executor.cores", "4") + .config("spark.master", "local") .getOrCreate() ) - session_name = f"projects/{test_project}/locations/{test_region}/sessions/{ManagedSparkSession._active_s8s_session_id}" + session_name = ( + f"projects/{test_project}/locations/{test_region}/sessions/" + f"{ManagedSparkSession._active_s8s_session_id}" + ) get_session_request = GetSessionRequest() get_session_request.name = session_name @@ -351,18 +248,21 @@ def test_create_spark_session_with_session_template_and_user_provided_session_co assert session.state == Session.State.ACTIVE assert session.session_template == session_template_name + assert ( + session.runtime_config.version + == ManagedSparkSession._DEFAULT_RUNTIME_VERSION + ) assert ( session.environment_config.execution_config.ttl == datetime.timedelta(seconds=64800) ) assert ( - session.runtime_config.properties["spark:spark.executor.cores"] == "16" + session.runtime_config.properties["spark:spark.executor.cores"] == "4" ) + assert session.runtime_config.properties["spark:spark.master"] == "local" assert ManagedSparkSession._active_s8s_session_uuid is not None connect_session.stop() - get_session_request = GetSessionRequest() - get_session_request.name = session_name session = session_controller_client.get_session(get_session_request) assert session.state in [ @@ -381,7 +281,7 @@ def test_add_artifacts_pypi_package(): Note: Skipped in CI due to infrastructure issues with PyPI package installation. """ connect_session = ManagedSparkSession.builder.getOrCreate() - from pyspark.sql.connect.functions import udf, sum + from pyspark.sql.connect.functions import sum, udf from pyspark.sql.types import IntegerType def generate_random2(row) -> int: @@ -391,7 +291,6 @@ def generate_random2(row) -> int: connect_session.addArtifacts("random2", pypi=True) - # Force evaluation of udf using random2 on workers sum_random = ( connect_session.range(1, 10) .withColumn( @@ -405,357 +304,203 @@ def generate_random2(row) -> int: connect_session.stop() -def test_sql_functions(connect_session): - """Test basic SQL functions like col(), sum(), count(), etc.""" - # Import SparkConnect-compatible functions - from pyspark.sql.connect.functions import col, sum, count - - # Create a test DataFrame - df = connect_session.createDataFrame( - [(1, "Alice", 100), (2, "Bob", 200), (3, "Charlie", 150)], - ["id", "name", "amount"], - ) - - # Test col() function - result_col = df.select(col("name")).collect() - assert len(result_col) == 3 - assert result_col[0]["name"] == "Alice" - - # Test aggregation functions - sum_result = df.select(sum("amount")).collect()[0][0] - assert sum_result == 450 - - count_result = df.select(count("id")).collect()[0][0] - assert count_result == 3 - - # Test with where clause using col() - filtered_df = df.where(col("amount") > 150) - filtered_count = filtered_df.count() - assert filtered_count == 1 - - # Test multiple column operations - df_with_calc = df.select( - col("id"), - col("name"), - col("amount"), - (col("amount") * 0.1).alias("tax"), - ) - tax_results = df_with_calc.collect() - assert tax_results[0]["tax"] == 10.0 - assert tax_results[1]["tax"] == 20.0 - assert tax_results[2]["tax"] == 15.0 - - -def test_sql_udf(connect_session): - """Test SQL UDF registration and usage.""" - # Import SparkConnect-compatible functions - from pyspark.sql.connect.functions import col, udf - - # Create a test DataFrame - df = connect_session.createDataFrame( - [(1, "hello"), (2, "world"), (3, "spark")], ["id", "text"] - ) - - # Register DataFrame for SQL queries - df.createOrReplaceTempView("test_table") - - # Define and register a Python UDF - def uppercase_func(text): - return text.upper() if text else None - - # Test UDF with DataFrame API - uppercase_udf = udf(uppercase_func, StringType()) - df_with_udf = df.select( - "id", "text", uppercase_udf(col("text")).alias("upper_text") - ) - df_result = df_with_udf.collect() - - assert df_result[0]["upper_text"] == "HELLO" - assert df_result[1]["upper_text"] == "WORLD" - - # Clean up - connect_session.sql("DROP VIEW IF EXISTS test_table") - - def test_session_reuse_with_custom_id( auth_type, test_project, test_region, session_controller_client, - os_environment, + isolated_session_state, ): - """Test the real-world session reuse scenario: create → terminate → recreate with same ID.""" - # Use a randomized session ID to avoid conflicts between test runs + """Test named session lifecycle: create -> reuse -> stop (stays active) -> reattach -> stop(terminate=True).""" custom_session_id = f"ml-pipeline-session-{uuid.uuid4().hex[:8]}" - - # Stop any existing session first to ensure clean state - if ManagedSparkSession._active_s8s_session_id: - try: - existing_session = ManagedSparkSession.getActiveSession() - if existing_session: - existing_session.stop() - except Exception: - pass - - # PHASE 1: Create initial session with custom ID - spark1 = ( - ManagedSparkSession.builder.sessionId(custom_session_id) - .projectId(test_project) - .location(test_region) - .getOrCreate() - ) - - # Verify session is created with custom ID - assert ManagedSparkSession._active_s8s_session_id == custom_session_id - first_session_uuid = spark1._active_s8s_session_uuid - - # Test basic functionality - df1 = spark1.createDataFrame([(1, "initial")], ["id", "stage"]) - result1 = df1.count() - assert result1 == 1 - - # PHASE 2: Test session reuse while active - # Clear cache to force session lookup - ManagedSparkSession._default_session = None - - spark2 = ( - ManagedSparkSession.builder.sessionId(custom_session_id) - .projectId(test_project) - .location(test_region) - .getOrCreate() - ) - - # Should reuse the same active session - assert ManagedSparkSession._active_s8s_session_id == custom_session_id - assert spark2._active_s8s_session_uuid == first_session_uuid - - # Test functionality on reused session - df2 = spark2.createDataFrame([(2, "reused")], ["id", "stage"]) - result2 = df2.count() - assert result2 == 1 - - # PHASE 3: Stop should not terminate named session - spark2.stop() - - # PHASE 4: Recreate with same ID - this tests the cleanup and recreation logic - # Clear all session state to ensure fresh lookup - ManagedSparkSession._default_session = None - ManagedSparkSession._active_s8s_session_id = None - ManagedSparkSession._active_s8s_session_uuid = None - - spark3 = ( - ManagedSparkSession.builder.sessionId(custom_session_id) - .projectId(test_project) - .location(test_region) - .getOrCreate() - ) - - # Should be a same session and same ID - assert ManagedSparkSession._active_s8s_session_id == custom_session_id - third_session_uuid = spark3._active_s8s_session_uuid - - # Should be same UUID - assert third_session_uuid == first_session_uuid - - # Test functionality on recreated session - df3 = spark3.createDataFrame([(3, "recreated")], ["id", "stage"]) - result3 = df3.count() - assert result3 == 1 - - # Clean up - spark3.stop() - - -def test_session_id_validation_in_integration( - test_project, test_region, os_environment -): - """Test session ID validation in integration environment.""" - - # Test invalid session ID raises ValueError - with pytest.raises(ValueError) as exc_info: - ManagedSparkSession.builder.sessionId("123-invalid-id") - assert "Invalid session ID" in str(exc_info.value) - - # Test that valid session ID works - valid_id = "valid-session-id-123" - builder = ( - ManagedSparkSession.builder.sessionId(valid_id) - .projectId(test_project) - .location(test_region) + session_name = ( + f"projects/{test_project}/locations/{test_region}/sessions/" + f"{custom_session_id}" ) - # Should not raise an exception - assert builder._custom_session_id == valid_id - + try: + # PHASE 1: Create initial session with custom ID in single-node mode + spark1 = ( + ManagedSparkSession.builder.sessionId(custom_session_id) + .projectId(test_project) + .location(test_region) + .config("spark.master", "local") + .getOrCreate() + ) -def test_sparksql_magic_with_managed_spark_session(connect_session): - """Test that %%sparksql works with the active ManagedSparkSession.""" - pytest.importorskip( - "IPython", reason="IPython not available (install [interactive])" - ) - pytest.importorskip( - "sqlparse", reason="sqlparse not available (install [interactive])" - ) + assert ManagedSparkSession._active_s8s_session_id == custom_session_id + first_session_uuid = spark1._active_s8s_session_uuid - from IPython.terminal.interactiveshell import TerminalInteractiveShell + df1 = spark1.createDataFrame([(1, "initial")], ["id", "stage"]) + assert df1.count() == 1 - # ManagedSparkSession is already the active session. - shell = TerminalInteractiveShell.instance() - shell.run_line_magic("load_ext", "google.cloud.managed_spark_magics") - shell.user_ns["factor"] = 5 - - shell.run_cell_magic( - "sparksql", - "-q result_df", - """ - SELECT - 10 * :factor as multiplication, - SQRT(16) as square_root, - CONCAT('Dataproc', '-', 'Spark') as joined_string - """, - ) + # PHASE 2: Test session reuse while active when _default_session is cleared + ManagedSparkSession._default_session = None - # Verify the result is captured in the namespace - assert "result_df" in shell.user_ns - df = shell.user_ns["result_df"] - assert df is not None + spark2 = ( + ManagedSparkSession.builder.sessionId(custom_session_id) + .projectId(test_project) + .location(test_region) + .getOrCreate() + ) - # Verify the computed values - data = df.collect() - assert len(data) == 1 - row = data[0] + assert ManagedSparkSession._active_s8s_session_id == custom_session_id + assert spark2._active_s8s_session_uuid == first_session_uuid - assert row["multiplication"] == 50 - assert row["square_root"] == 4.0 - assert row["joined_string"] == "Dataproc-Spark" + df2 = spark2.createDataFrame([(2, "reused")], ["id", "stage"]) + assert df2.count() == 1 + # PHASE 3: Default stop() on a named session does NOT terminate on the server + spark2.stop() + assert ManagedSparkSession._active_s8s_session_id is None -def test_stop_named_session_with_terminate_true( - auth_type, - test_project, - test_region, - session_controller_client, - os_environment, -): - """Test that stop(terminate=True) terminates a named session on the server.""" - # Use a randomized session ID to avoid conflicts - custom_session_id = f"test-terminate-true-{uuid.uuid4().hex[:8]}" - - # Create a session with custom ID - spark = ( - ManagedSparkSession.builder.sessionId(custom_session_id) - .projectId(test_project) - .location(test_region) - .getOrCreate() - ) + get_session_request = GetSessionRequest(name=session_name) + session = session_controller_client.get_session(get_session_request) + assert session.state == Session.State.ACTIVE - # Verify session is created - assert ManagedSparkSession._active_s8s_session_id == custom_session_id - session_name = f"projects/{test_project}/locations/{test_region}/sessions/{custom_session_id}" + # PHASE 4: Reattach to the active server-side session with the same custom ID + ManagedSparkSession._default_session = None + ManagedSparkSession._active_s8s_session_id = None + ManagedSparkSession._active_s8s_session_uuid = None - # Test basic functionality - df = spark.createDataFrame([(1, "test")], ["id", "value"]) - assert df.count() == 1 + spark3 = ( + ManagedSparkSession.builder.sessionId(custom_session_id) + .projectId(test_project) + .location(test_region) + .getOrCreate() + ) - # Stop with terminate=True - spark.stop(terminate=True) + assert ManagedSparkSession._active_s8s_session_id == custom_session_id + assert spark3._active_s8s_session_uuid == first_session_uuid - # Verify client-side cleanup - assert ManagedSparkSession._active_s8s_session_id is None + df3 = spark3.createDataFrame([(3, "recreated")], ["id", "stage"]) + assert df3.count() == 1 - # Verify server-side session is terminating or terminated - get_session_request = GetSessionRequest() - get_session_request.name = session_name - session = session_controller_client.get_session(get_session_request) + # PHASE 5: stop(terminate=True) terminates the named session on the server + spark3.stop(terminate=True) + assert ManagedSparkSession._active_s8s_session_id is None - assert session.state in [ - Session.State.TERMINATING, - Session.State.TERMINATED, - ] + session = session_controller_client.get_session(get_session_request) + assert session.state in [ + Session.State.TERMINATING, + Session.State.TERMINATED, + ] + finally: + # Ensure the named session is never leaked on the server if an assertion fails mid-test + try: + session = session_controller_client.get_session( + GetSessionRequest(name=session_name) + ) + if session.state in [ + Session.State.CREATING, + Session.State.ACTIVE, + ]: + session_controller_client.terminate_session( + TerminateSessionRequest(name=session_name) + ) + except Exception: + pass -def test_stop_managed_session_with_terminate_false( - auth_type, +def test_serverless_session_termination_and_deletion_lifecycle( test_project, test_region, session_controller_client, - os_environment, + isolated_session_state, ): - """Test that stop(terminate=False) does NOT terminate a managed session on the server.""" - # Create a managed session (auto-generated ID) - spark = ( - ManagedSparkSession.builder.projectId(test_project) - .location(test_region) - .getOrCreate() - ) - - # Verify it's a managed session (auto-generated ID) - assert ManagedSparkSession._active_s8s_session_id is not None - assert ManagedSparkSession._active_session_uses_custom_id is False - session_id = ManagedSparkSession._active_s8s_session_id - session_name = ( - f"projects/{test_project}/locations/{test_region}/sessions/{session_id}" - ) - - # Test basic functionality - df = spark.createDataFrame([(1, "test")], ["id", "value"]) - assert df.count() == 1 - - # Stop with terminate=False (prevent server-side termination) - spark.stop(terminate=False) - - # Verify client-side cleanup - assert ManagedSparkSession._active_s8s_session_id is None - - # Verify server-side session is still ACTIVE (not terminated) - get_session_request = GetSessionRequest() - get_session_request.name = session_name - session = session_controller_client.get_session(get_session_request) - - assert session.state == Session.State.ACTIVE - - # Clean up: terminate the session manually - terminate_session_request = TerminateSessionRequest() - terminate_session_request.name = session_name - session_controller_client.terminate_session(terminate_session_request) - - -@pytest.fixture -def batch_workload_env(monkeypatch): - """Sets DATAPROC_WORKLOAD_TYPE to 'batch' for a test.""" - monkeypatch.setenv("DATAPROC_WORKLOAD_TYPE", "batch") - + """Test managed session stop(terminate=False), server-side termination recovery, and server-side deletion.""" + first_session_name = None + second_session_name = None + try: + # PHASE 1: Create a managed session (auto-generated ID) in single-node mode + # and verify stop(terminate=False) leaves the server-side session ACTIVE. + first_spark = ( + ManagedSparkSession.builder.projectId(test_project) + .location(test_region) + .config("spark.master", "local") + .getOrCreate() + ) + assert ManagedSparkSession._active_s8s_session_id is not None + assert ManagedSparkSession._active_session_uses_custom_id is False + first_session_id = ManagedSparkSession._active_s8s_session_id + first_session_uuid = ManagedSparkSession._active_s8s_session_uuid + first_session_name = ( + f"projects/{test_project}/locations/{test_region}/sessions/" + f"{first_session_id}" + ) -@pytest.fixture -def local_spark_session(): - """Provides a standard local PySpark session for comparison.""" - from pyspark.sql import SparkSession as PySparkSession + df = first_spark.createDataFrame([(1, "test")], ["id", "value"]) + assert df.count() == 1 - # Stop any existing session to ensure a clean environment for creating a local session. - # This prevents test isolation failures where a Managed Spark session from a previous - # test might be picked up by getOrCreate(). - if ManagedSparkSession.getActiveSession(): - ManagedSparkSession.getActiveSession().stop() + # Stop client without terminating server + first_spark.stop(terminate=False) + assert ManagedSparkSession._active_s8s_session_id is None - session = PySparkSession.builder.master("local").getOrCreate() - yield session - session.stop() + get_first_request = GetSessionRequest(name=first_session_name) + session = session_controller_client.get_session(get_first_request) + assert session.state == Session.State.ACTIVE + # PHASE 2: Re-associate client with first_session, terminate it on the + # server, and verify getOrCreate() stops the terminated session and creates + # a replacement session. + ManagedSparkSession._active_s8s_session_id = first_session_id + ManagedSparkSession._active_s8s_session_uuid = first_session_uuid -@pytest.mark.skipif( - is_remote_only(), - reason="Local SparkSession requires full pyspark (not pyspark-client).", -) -def test_create_local_spark_session(batch_workload_env, local_spark_session): - """Test creating a local Spark session.""" - from pyspark.sql import SparkSession as PySparkSession + terminate_request = TerminateSessionRequest(name=first_session_name) + session_controller_client.terminate_session(terminate_request).result() - managed_spark_session = ManagedSparkSession.builder.getOrCreate() - try: - assert isinstance(managed_spark_session, PySparkSession) - assert not isinstance(managed_spark_session, ManagedSparkSession) + second_spark = ( + ManagedSparkSession.builder.projectId(test_project) + .location(test_region) + .config("spark.master", "local") + .getOrCreate() + ) + second_session_id = ManagedSparkSession._active_s8s_session_id + second_session_uuid = ManagedSparkSession._active_s8s_session_uuid + second_session_name = ( + f"projects/{test_project}/locations/{test_region}/sessions/" + f"{second_session_id}" + ) - # Compare configurations to ensure they are both local sessions - assert managed_spark_session == local_spark_session + assert first_session_uuid != second_session_uuid + assert second_session_uuid is not None + assert second_session_id is not None + + first_state = session_controller_client.get_session( + get_first_request + ).state + assert first_state in [ + Session.State.TERMINATING, + Session.State.TERMINATED, + ] + + get_second_request = GetSessionRequest(name=second_session_name) + second_state = session_controller_client.get_session( + get_second_request + ).state + assert second_state == Session.State.ACTIVE + + # PHASE 3: Delete the active session on the server and verify stop() cleans + # up client state without error. + delete_request = DeleteSessionRequest(name=second_session_name) + session_controller_client.delete_session(delete_request).result() + + second_spark.stop() + assert ManagedSparkSession._active_s8s_session_uuid is None + assert ManagedSparkSession._active_s8s_session_id is None finally: - managed_spark_session.stop() + for name in (first_session_name, second_session_name): + if name: + try: + session = session_controller_client.get_session( + GetSessionRequest(name=name) + ) + if session.state in [ + Session.State.CREATING, + Session.State.ACTIVE, + ]: + session_controller_client.terminate_session( + TerminateSessionRequest(name=name) + ) + except Exception: + pass diff --git a/tests/local/test_sparksql_local.py b/tests/local/test_sparksql_local.py index a004804..90f820f 100644 --- a/tests/local/test_sparksql_local.py +++ b/tests/local/test_sparksql_local.py @@ -78,6 +78,12 @@ def _java_version(): @pytest.fixture(scope="module", params=["classic", "connect"]) def spark(request): with pytest.MonkeyPatch.context() as mp: + mp.setenv("SPARK_LOCAL_IP", "127.0.0.1") + mp.setenv("NO_PROXY", "localhost,127.0.0.1") + mp.setenv("no_proxy", "localhost,127.0.0.1") + mp.delenv("SPARK_LOCAL_REMOTE", raising=False) + mp.delenv("SPARK_CONNECT_MODE_ENABLED", raising=False) + mp.delenv("SPARK_CONNECT_AUTHENTICATE_TOKEN", raising=False) builder = SparkSession.builder if request.param == "classic": builder = builder.master("local[2]").config( @@ -97,6 +103,8 @@ def spark(request): ) yield session session.stop() + if SparkSession._instantiatedSession is not None: + SparkSession._instantiatedSession.stop() @pytest.fixture @@ -303,3 +311,27 @@ def test_full_traceback(ip): sparksql.get_config(ip).full_traceback = True with pytest.raises(PySparkException): ip.run_cell_magic("sparksql", "", "SELECT nope") + + +def test_create_local_spark_session(monkeypatch): + """Test creating a local Spark session when DATAPROC_WORKLOAD_TYPE is batch.""" + from google.cloud.managed_spark_connect import ManagedSparkSession + + monkeypatch.setenv("DATAPROC_WORKLOAD_TYPE", "batch") + monkeypatch.setenv("SPARK_LOCAL_IP", "127.0.0.1") + + if ManagedSparkSession.getActiveSession(): + ManagedSparkSession.getActiveSession().stop() + + local_spark_session = ( + SparkSession.builder.master("local") + .config("spark.ui.enabled", "false") + .getOrCreate() + ) + try: + managed_spark_session = ManagedSparkSession.builder.getOrCreate() + assert isinstance(managed_spark_session, SparkSession) + assert not isinstance(managed_spark_session, ManagedSparkSession) + assert managed_spark_session == local_spark_session + finally: + local_spark_session.stop() diff --git a/tests/unit/test_init.py b/tests/unit/test_init.py index ea92d8e..ab57ec8 100644 --- a/tests/unit/test_init.py +++ b/tests/unit/test_init.py @@ -1009,17 +1009,65 @@ def test_all_import_orders_load_all_extras(self): _NO_INTERACTIVE_DEPS_PROBE = """ import sys +from unittest import mock # Simulate an install without the [interactive] extra. A None entry makes both # importlib.util.find_spec() and import treat the module as absent. -for name in ("IPython", "sqlparse", "traitlets", "google.colabsqlviz"): +for name in ( + "IPython", + "sqlparse", + "traitlets", + "google.colabsqlviz", + "ipywidgets", +): sys.modules[name] = None import google.cloud.managed_spark_connect as msc +from google.cloud.dataproc_v1 import Session +from pyspark.sql.connect.client.core import ConfigResult +from pyspark.sql.connect.proto import ConfigResponse print(msc.ManagedSparkSession is not None) print(hasattr(msc, "ManagedSparkConnect")) print("google.cloud.managed_spark_connect._ipython" in sys.modules) + +# Also verify ManagedSparkSession creation, tqdm fallback, HTML repr, and stop +# work without any [interactive] dependencies present. +with ( + mock.patch("google.auth.default") as mock_creds, + mock.patch("google.cloud.dataproc_v1.SessionControllerClient") as mock_client_cls, + mock.patch("pyspark.sql.connect.client.SparkConnectClient.config") as mock_config, + mock.patch("google.cloud.managed_spark_connect.session.is_s8s_session_active", return_value=True), +): + cred = mock.MagicMock() + cred.token = "token" + mock_creds.return_value = (cred, "") + mock_config.return_value = ConfigResult.fromProto(ConfigResponse()) + + session_response = Session() + session_response.runtime_info.endpoints["Spark Connect Server"] = ( + "sc://spark-connect-server.example.com:443" + ) + session_response.uuid = "c002e4ef-fe5e-41a8-a157-160aa73e4f7f" + mock_op = mock.Mock() + mock_op.result.return_value = session_response + mock_client_Instance = mock_client_cls.return_value + mock_client_Instance.create_session.return_value = mock_op + term_response = Session() + term_response.state = Session.State.TERMINATING + mock_client_Instance.get_session.return_value = term_response + + session = ( + msc.ManagedSparkSession.builder.projectId("test-project") + .location("test-region") + .getOrCreate() + ) + try: + print(isinstance(session, msc.ManagedSparkSession)) + print(msc.ManagedSparkSession.get_tqdm_bar() is not None) + print("Spark Connect" in session._repr_html_()) + finally: + session.stop() """ @@ -1044,7 +1092,10 @@ def test_import_without_interactive_deps(self): env=env, ) self.assertEqual(result.returncode, 0, result.stderr) - self.assertEqual(result.stdout.split(), ["True", "False", "False"]) + self.assertEqual( + result.stdout.split(), + ["True", "False", "False", "True", "True", "True"], + ) _NO_PYSPARK_PROBE = """ diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index 46814a6..5784ef7 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -22,6 +22,7 @@ InvalidArgument, NotFound, ) +from google.auth.exceptions import DefaultCredentialsError from google.cloud.managed_spark_connect import ManagedSparkSession from google.cloud.managed_spark_connect.exceptions import ManagedSparkConnectException from google.cloud.managed_spark_connect.session import ( @@ -1603,30 +1604,34 @@ def test_create_session_with_valid_notebook_id( def test_create_session_without_project_id(self): """Tests that an exception is raised when project ID is not provided.""" os.environ.clear() - try: + with self.assertRaises(ManagedSparkConnectException) as context: ManagedSparkSession.builder.location("test-region").getOrCreate() - except ManagedSparkConnectException as e: - self.assertIn("project ID is not set", str(e)) + self.assertIn("project ID is not set", str(context.exception)) def test_create_session_without_location(self): """Tests that an exception is raised when location is not provided.""" os.environ.clear() - try: + with self.assertRaises(ManagedSparkConnectException) as context: ManagedSparkSession.builder.projectId("test-project").getOrCreate() - except ManagedSparkConnectException as e: - self.assertIn("location is not set", str(e)) + self.assertIn("location is not set", str(context.exception)) - def test_create_session_without_application_default_credentials(self): + @mock.patch( + "google.auth.default", + side_effect=DefaultCredentialsError("No ADC"), + ) + def test_create_session_without_application_default_credentials( + self, mock_default_creds + ): """Tests that an exception is raised when application default credentials is not provided.""" os.environ.clear() - try: + with self.assertRaises(ManagedSparkConnectException) as context: ManagedSparkSession.builder.location("test-region").projectId( "test-project" ).getOrCreate() - except ManagedSparkConnectException as e: - self.assertIn( - "Credentials error while creating Managed Spark Session", str(e) - ) + self.assertIn( + "Credentials error while creating Managed Spark Session", + str(context.exception), + ) class ManagedSparkConnectClientTest(unittest.TestCase): @@ -2729,7 +2734,11 @@ def test_dataproc_session_id_builder_method(self): builder = ManagedSparkSession.builder # Test valid session ID - result = builder.sessionId("test-session") + result = ( + builder.sessionId("test-session") + .projectId("test-project") + .location("test-region") + ) self.assertEqual(builder._custom_session_id, "test-session") self.assertEqual(result, builder) # Check method chaining