diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml index 4205d67c..cb7bc98c 100644 --- a/.github/workflows/tests.yaml +++ b/.github/workflows/tests.yaml @@ -28,7 +28,7 @@ concurrency: jobs: test: - name: Run unit tests (${{ matrix.os }}, [${{ matrix.spark-extra }}]) + name: Run unit tests (${{ matrix.os }}, Python ${{ matrix.python-version }}, [${{ matrix.spark-extra }}]${{ matrix.smoke && ', smoke' || '' }}) runs-on: ${{ matrix.os }} timeout-minutes: 15 defaults: @@ -38,9 +38,33 @@ jobs: fail-fast: false matrix: os: ["ubuntu-latest", "macos-latest", "windows-latest"] + # The Python version of the default Managed Spark runtime, which + # matters most since Python UDFs need client and server to match. + python-version: ["3.12"] # Mutually exclusive: [client] installs pyspark-client, [spark] full # pyspark[connect]. spark-extra: ["client", "spark"] + # The rest of the versions setup.py declares get smoke tests only, on + # one OS and extra: they are there to catch syntax, dependency + # resolution, import and standard library problems, which are neither + # OS- nor extra-specific. + include: + - os: "ubuntu-latest" + python-version: "3.10" + spark-extra: "client" + smoke: true + - os: "ubuntu-latest" + python-version: "3.11" + spark-extra: "client" + smoke: true + - os: "ubuntu-latest" + python-version: "3.13" + spark-extra: "client" + smoke: true + - os: "ubuntu-latest" + python-version: "3.14" + spark-extra: "client" + smoke: true steps: - name: Checkout code @@ -51,7 +75,7 @@ jobs: - name: Setup Python uses: actions/setup-python@6d4fd22b20465ccc0f7354476588e8df7c9ca0e0 # v7.1.0 with: - python-version: "3.12" + python-version: ${{ matrix.python-version }} - name: Set up uv uses: astral-sh/setup-uv@18d369b23da25a81e9458561248e1aafff6a70fa # v10.3.1 @@ -63,10 +87,10 @@ jobs: ~/.cache/uv ~/Library/Caches/uv ~/AppData/Local/uv/cache - key: ${{ runner.os }}-uv-${{ matrix.spark-extra }}-${{ hashFiles('setup.py', 'requirements-test.txt') }} + key: ${{ runner.os }}-uv-py${{ matrix.python-version }}-${{ matrix.spark-extra }}-${{ hashFiles('setup.py', 'requirements-test.txt') }} restore-keys: | - ${{ runner.os }}-uv-${{ matrix.spark-extra }}- - ${{ runner.os }}-uv- + ${{ runner.os }}-uv-py${{ matrix.python-version }}-${{ matrix.spark-extra }}- + ${{ runner.os }}-uv-py${{ matrix.python-version }}- # Installs through the package's own extras, rather than # requirements-dev.txt (which pins pyspark-client), so that both extras are @@ -76,14 +100,24 @@ jobs: 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 + run: python -m pytest tests/unit/test_init.py tests/unit/test_pypi_artifacts.py -k "TestWithoutInteractiveExtra or TestPythonVersionCheck or TestPySparkVersionCheck 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 + if: ${{ !matrix.smoke }} run: python -m pytest tests/unit/ -v --tb=short -n auto + # All but the slow unit tests. Those are mostly waiting, and cover next to + # nothing the others do not, so this covers nearly every line the full suite + # does in a fraction of the time, catching standard library APIs this Python + # lacks as well as syntax. (The tests import every module; vermin.yaml + # checks syntax statically everywhere else.) + - name: Run smoke tests + if: ${{ matrix.smoke }} + run: python -m pytest tests/unit/ -m "not slow" -v --tb=short -n auto + # Starts Spark in-process, so it runs serially, outside of xdist. Needs # the JVM that only full pyspark ships. - name: Run local Spark tests @@ -91,3 +125,42 @@ jobs: run: | export JAVA_HOME="${JAVA_HOME_17_X64:-${JAVA_HOME_17_ARM64:-${JAVA_HOME_21_X64:-${JAVA_HOME_21_ARM64:-$JAVA_HOME}}}}" python -m pytest tests/local/ -v --tb=short + + # The [client] and [spark] extras pin the pyspark minor version of the + # default Managed Spark runtime. Users of other runtimes, and environments + # that preinstall pyspark, can have a newer one, so also run the full suite + # against the newest pyspark-client. On Python 3.13 rather than 3.12, which + # the main matrix already covers with the pinned pyspark. + test-newest-pyspark: + name: Run unit tests (ubuntu-latest, Python 3.13, newest pyspark-client) + runs-on: ubuntu-latest + timeout-minutes: 15 + defaults: + run: + shell: bash + + steps: + - name: Checkout code + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + persist-credentials: false + + - name: Setup Python + uses: actions/setup-python@6d4fd22b20465ccc0f7354476588e8df7c9ca0e0 # v7.1.0 + with: + python-version: "3.13" + + - name: Set up uv + uses: astral-sh/setup-uv@18d369b23da25a81e9458561248e1aafff6a70fa # v10.3.1 + + # No extra: it would pin pyspark-client back to the default runtime's. + - name: Install dependencies + run: uv pip install --system -e '.[interactive]' pyspark-client -r requirements-test.txt + + # Serially, without xdist: pyspark 4.2 exits on import when __main__ has + # no __file__, as in xdist workers, mistaking that for running its own + # doctests. + - name: Run unit tests + run: | + python -c 'import pyspark; print("pyspark", pyspark.__version__)' + python -m pytest tests/unit/ -v --tb=short diff --git a/.github/workflows/vermin.yaml b/.github/workflows/vermin.yaml new file mode 100644 index 00000000..3ea7ca26 --- /dev/null +++ b/.github/workflows/vermin.yaml @@ -0,0 +1,62 @@ +# 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. + +# Statically checks that the code only uses language features and standard +# library modules, functions and arguments that are available in the oldest +# Python version setup.py declares. Complements the smoke tests in tests.yaml, +# which only catch the APIs they run, and are run on CI's Python versions only. +# vermin only knows the types of names it can resolve statically, so it can +# miss methods on objects whose types it cannot infer. Silence false positives +# with a "# novm" comment. +name: Vermin +on: + push: + branches: [ main ] + pull_request: + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + +jobs: + vermin: + name: Check minimum Python version + runs-on: ubuntu-latest + timeout-minutes: 5 + + steps: + - name: Checkout code + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + persist-credentials: false + + - name: Setup Python + uses: actions/setup-python@6d4fd22b20465ccc0f7354476588e8df7c9ca0e0 # v7.1.0 + with: + python-version: "3.12" + + - name: Set up uv + uses: astral-sh/setup-uv@18d369b23da25a81e9458561248e1aafff6a70fa # v10.3.1 + + - name: Install vermin + run: uv pip install --system 'vermin~=1.9' + + # The target must match python_requires in setup.py. The trailing "-" fails + # only if the code needs a newer version, not if it would also run on an + # older one. + - name: Run vermin + run: vermin --no-tips --violations --eval-annotations --target=3.10- google tests setup.py diff --git a/README.md b/README.md index 97508091..e4bf8907 100644 --- a/README.md +++ b/README.md @@ -35,6 +35,13 @@ Other installation options: > later uninstalled. In particular, never combine the `client` and `spark` > extras (`'[client,spark]'`). +### Python versions + +This package requires Python 3.10 or later. To use Python UDFs (including +`applyInPandas`, `mapInPandas`, and similar APIs), the client's Python minor +version must match the Python of the Managed Spark runtime the session runs on, +for example Python 3.12 for runtime 3.0. The client warns when they differ. + ## Uninstall ```sh diff --git a/google/cloud/managed_spark_connect/_runtimes.py b/google/cloud/managed_spark_connect/_runtimes.py new file mode 100644 index 00000000..7163bc85 --- /dev/null +++ b/google/cloud/managed_spark_connect/_runtimes.py @@ -0,0 +1,73 @@ +# 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. +"""What the Managed Spark runtime versions ship, for client compatibility checks. + +Spark Connect needs the client's PySpark minor version to match the server's +Spark minor version, and Python UDFs need the client's Python minor version to +match the server's, since PySpark refuses to run a UDF pickled by a different +Python minor version. +""" + +import dataclasses +from typing import Optional, Tuple + +from packaging.version import InvalidVersion, Version + +MajorMinor = Tuple[int, int] + + +@dataclasses.dataclass(frozen=True) +class RuntimeInfo: + """The versions a Managed Spark runtime version ships. + + Attributes: + spark: Apache Spark (major, minor) version. + python: Python (major, minor) version. + """ + + spark: MajorMinor + python: MajorMinor + + +# Keyed by runtime (major, minor) version. Runtime versions missing from here +# are not checked, so a runtime newer than this release still works, just +# without the compatibility warnings. +_RUNTIMES = { + (3, 0): RuntimeInfo(spark=(4, 0), python=(3, 12)), +} + + +def major_minor(version_str: str) -> Optional[MajorMinor]: + """The (major, minor) of a version string, or None if it has no minor. + + Accepts both runtime versions ("3.0", or subminor versions like "3.0.12") + and PySpark versions ("4.0.1", "4.1.0.dev1"). + """ + try: + release = Version(version_str).release + except InvalidVersion: + return None + if len(release) < 2: + return None + return release[0], release[1] + + +def get_runtime_info(runtime_version: str) -> Optional[RuntimeInfo]: + """What the given runtime version ships, or None if unknown.""" + key = major_minor(runtime_version) + return _RUNTIMES.get(key) if key else None + + +def format_version(version: MajorMinor) -> str: + return f"{version[0]}.{version[1]}" diff --git a/google/cloud/managed_spark_connect/session.py b/google/cloud/managed_spark_connect/session.py index e00a3dc1..d38955cc 100644 --- a/google/cloud/managed_spark_connect/session.py +++ b/google/cloud/managed_spark_connect/session.py @@ -85,6 +85,7 @@ TerminateSessionRequest, ) from google.cloud.dataproc_v1.types import sessions +from google.cloud.managed_spark_connect import _runtimes from google.cloud.managed_spark_connect import environment from google.cloud.managed_spark_connect import ide from pyspark.sql.connect.session import SparkSession @@ -1185,6 +1186,9 @@ def _get_session_config(self): self._check_python_version_compatibility( session_config.runtime_config.version ) + self._check_pyspark_version_compatibility( + session_config.runtime_config.version + ) # Use local variable to improve readability of deeply nested attribute access exec_config = session_config.environment_config.execution_config @@ -1365,24 +1369,57 @@ def _check_python_version_compatibility(self, runtime_version): import sys import warnings - # Runtime version to server Python version mapping - RUNTIME_PYTHON_MAP = { - "3.0": (3, 12), - } + runtime_info = _runtimes.get_runtime_info(runtime_version) + if runtime_info is None: + return client_python = sys.version_info[:2] # (major, minor) + server_python = runtime_info.python + + if client_python != server_python: + warnings.warn( + f"Python version mismatch detected: Client is using Python {client_python[0]}.{client_python[1]}, " + f"but Managed Spark runtime {runtime_version} uses Python {server_python[0]}.{server_python[1]}. " + f"This mismatch may cause issues with Python UDF (User Defined Function) compatibility. " + f"Consider using Python {server_python[0]}.{server_python[1]} for optimal UDF execution.", + stacklevel=3, + ) - if runtime_version in RUNTIME_PYTHON_MAP: - server_python = RUNTIME_PYTHON_MAP[runtime_version] + def _check_pyspark_version_compatibility(self, runtime_version): + """Check that the installed PySpark is not newer than the runtime's Spark. - if client_python != server_python: - warnings.warn( - f"Python version mismatch detected: Client is using Python {client_python[0]}.{client_python[1]}, " - f"but Managed Spark runtime {runtime_version} uses Python {server_python[0]}.{server_python[1]}. " - f"This mismatch may cause issues with Python UDF (User Defined Function) compatibility. " - f"Consider using Python {server_python[0]}.{server_python[1]} for optimal UDF execution.", - stacklevel=3, - ) + Spark Connect servers support clients with older Spark minor + versions, but not newer ones: for example, PySpark 4.2's + createDataFrame() reads SQL configs Spark 4.0 does not have. + """ + import warnings + + from pyspark.version import __version__ as pyspark_version + + runtime_info = _runtimes.get_runtime_info(runtime_version) + client_spark = _runtimes.major_minor(pyspark_version) + if runtime_info is None or client_spark is None: + return + + server_spark = runtime_info.spark + if client_spark > server_spark: + from pyspark.util import is_remote_only + + server_spark_str = _runtimes.format_version(server_spark) + # Downgrade whichever distribution is installed, so as not to + # end up with both writing to the pyspark module. + package = ( + "pyspark-client" if is_remote_only() else "pyspark[connect]" + ) + warnings.warn( + f"PySpark {pyspark_version} is installed, but Managed Spark " + f"runtime {runtime_version} runs Spark {server_spark_str}. " + "Spark Connect does not support clients newer than the " + "server, so some operations will fail. Install a matching " + f"PySpark, e.g. pip install '{package}~={server_spark_str}.0', " + "or choose a newer runtime version with runtimeVersion().", + stacklevel=3, + ) def _check_runtime_compatibility(self, session_config): """Check if runtime version 3.0 client is compatible with older runtime versions. diff --git a/pyproject.toml b/pyproject.toml index 56261e7d..521f88c5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,4 +10,5 @@ pyink-indentation = 4 # (default is 4) markers = [ "integration: marks tests as integration tests", "ci_safe: marks tests that work in CI environment", + "slow: marks unit tests that take a second or more, mostly waiting; the smoke tests skip them", ] diff --git a/setup.py b/setup.py index 77ad3926..f511a4e8 100644 --- a/setup.py +++ b/setup.py @@ -76,6 +76,23 @@ def run(self) -> None: author="Google LLC", url="https://github.com/GoogleCloudDataproc/managed-spark-connect-python", license="Apache 2.0", + # The floor is what the code needs (match statements), which is also what + # pyspark>=4.1 needs. Python UDFs additionally need the client's Python + # minor version to match the Managed Spark runtime's; the session warns + # about a mismatch rather than refusing to install, since everything else + # works across Python versions. When changing the floor, also change the + # vermin target in .github/workflows/vermin.yaml and the smoke tests in + # .github/workflows/tests.yaml. + python_requires=">=3.10", + classifiers=[ + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3 :: Only", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", + ], packages=find_namespace_packages(include=["google.*"]), exclude_package_data={"*": [".wh.*"]}, cmdclass={"egg_info": _CleanEggInfo}, diff --git a/tests/unit/test_init.py b/tests/unit/test_init.py index 939b1922..71181589 100644 --- a/tests/unit/test_init.py +++ b/tests/unit/test_init.py @@ -11,13 +11,33 @@ # 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 subprocess +import sys import unittest from unittest import mock +from google.cloud.managed_spark_connect import _runtimes from google.cloud.managed_spark_connect.session import ManagedSparkSession from google.cloud.managed_spark_connect.exceptions import ManagedSparkConnectException +def _run_python(code, *args, env=None): + """Runs code in a fresh interpreter, like `python -c code args...`. + + Feeds the code on stdin rather than with -c: pyspark 4.2 takes a __main__ + without __file__, as -c leaves it, to mean that it is running its own + doctests, and then exits when it cannot find a full Spark distribution, + which pyspark-client does not have. + """ + return subprocess.run( + [sys.executable, "-", *args], + input=code, + capture_output=True, + text=True, + env=env, + ) + + class TestPythonVersionCheck(unittest.TestCase): def test_python_version_mismatch_warning_for_runtime_30(self): @@ -69,6 +89,85 @@ def test_no_warning_for_unknown_runtime_version(self): mock_warn.assert_not_called() + def test_python_version_mismatch_warning_for_subminor_runtime(self): + """Subminor runtime versions are looked up by their major.minor.""" + with mock.patch("sys.version_info", (3, 11, 0)): + with mock.patch("warnings.warn") as mock_warn: + session_builder = ManagedSparkSession.Builder() + session_builder._check_python_version_compatibility("3.0.12") + + mock_warn.assert_called_once() + self.assertIn( + "Managed Spark runtime 3.0.12 uses Python 3.12", + mock_warn.call_args.args[0], + ) + + +class TestPySparkVersionCheck(unittest.TestCase): + + def _check(self, runtime_version, pyspark_version, remote_only=True): + with ( + mock.patch("pyspark.version.__version__", pyspark_version), + mock.patch("pyspark.util.is_remote_only", return_value=remote_only), + mock.patch("warnings.warn") as mock_warn, + ): + ManagedSparkSession.Builder()._check_pyspark_version_compatibility( + runtime_version + ) + return mock_warn + + def test_no_warning_when_spark_versions_match(self): + for pyspark_version in ["4.0.0", "4.0.5", "4.0.1.dev1"]: + with self.subTest(pyspark_version=pyspark_version): + self._check("3.0", pyspark_version).assert_not_called() + + def test_no_warning_when_pyspark_older_than_runtime(self): + """Spark Connect servers support older clients.""" + newer_runtime = { + (9, 0): _runtimes.RuntimeInfo(spark=(4, 2), python=(3, 13)) + } + with mock.patch.dict(_runtimes._RUNTIMES, newer_runtime): + for pyspark_version in ["4.0.0", "4.1.3"]: + with self.subTest(pyspark_version=pyspark_version): + self._check("9.0", pyspark_version).assert_not_called() + + def test_warning_when_pyspark_newer_than_runtime(self): + mock_warn = self._check("3.0", "4.1.2") + + mock_warn.assert_called_once_with( + "PySpark 4.1.2 is installed, but Managed Spark runtime 3.0 runs " + "Spark 4.0. Spark Connect does not support clients newer than the " + "server, so some operations will fail. Install a matching PySpark, " + "e.g. pip install 'pyspark-client~=4.0.0', or choose a newer " + "runtime version with runtimeVersion().", + stacklevel=3, + ) + + def test_warning_recommends_full_pyspark_when_installed(self): + mock_warn = self._check("3.0", "4.1.2", remote_only=False) + + mock_warn.assert_called_once() + self.assertIn( + "pip install 'pyspark[connect]~=4.0.0'", mock_warn.call_args.args[0] + ) + + def test_warning_for_subminor_runtime(self): + mock_warn = self._check("3.0.12", "4.1.2") + + mock_warn.assert_called_once() + self.assertIn( + "Managed Spark runtime 3.0.12 runs Spark 4.0", + mock_warn.call_args.args[0], + ) + + def test_no_warning_for_unknown_runtime_version(self): + for runtime_version in ["99.0", "unknown", ""]: + with self.subTest(runtime_version=runtime_version): + self._check(runtime_version, "4.1.2").assert_not_called() + + def test_no_warning_for_unparseable_pyspark_version(self): + self._check("3.0", "custom-build").assert_not_called() + class TestRuntimeVersionCompatibility(unittest.TestCase): @@ -992,17 +1091,10 @@ class TestExtrasImportOrder(unittest.TestCase): def _run(self, order): import os - import subprocess - import sys env = dict(os.environ) env.pop("MANAGED_SPARK_CONNECT_ENABLE_EXTRAS", None) - result = subprocess.run( - [sys.executable, "-c", _IMPORT_ORDER_PROBE, order], - capture_output=True, - text=True, - env=env, - ) + result = _run_python(_IMPORT_ORDER_PROBE, order, env=env) self.assertEqual(result.returncode, 0, result.stderr) return result.stdout.split() @@ -1092,17 +1184,10 @@ class TestWithoutInteractiveExtra(unittest.TestCase): def test_import_without_interactive_deps(self): import os - import subprocess - import sys env = dict(os.environ) env.pop("MANAGED_SPARK_CONNECT_ENABLE_EXTRAS", None) - result = subprocess.run( - [sys.executable, "-c", _NO_INTERACTIVE_DEPS_PROBE], - capture_output=True, - text=True, - env=env, - ) + result = _run_python(_NO_INTERACTIVE_DEPS_PROBE, env=env) self.assertEqual(result.returncode, 0, result.stderr) self.assertEqual( result.stdout.split(), @@ -1111,8 +1196,6 @@ def test_import_without_interactive_deps(self): def test_non_ipython_import_skips_heavy_optional_modules(self): import os - import subprocess - import sys probe = """ import sys @@ -1124,12 +1207,7 @@ def test_non_ipython_import_skips_heavy_optional_modules(self): """ env = dict(os.environ) env.pop("MANAGED_SPARK_CONNECT_ENABLE_EXTRAS", None) - result = subprocess.run( - [sys.executable, "-c", probe], - capture_output=True, - text=True, - env=env, - ) + result = _run_python(probe, env=env) self.assertEqual(result.returncode, 0, result.stderr) self.assertEqual( result.stdout.split(), @@ -1192,8 +1270,6 @@ def _run(self, ipython_mode, pyspark_mode, pyspark_version=None): import json import os import shutil - import subprocess - import sys import tempfile tmp = tempfile.mkdtemp() @@ -1207,12 +1283,8 @@ def _run(self, ipython_mode, pyspark_mode, pyspark_version=None): env = dict(os.environ) env["IPYTHONDIR"] = os.path.join(tmp, "ipython") - result = subprocess.run( - [sys.executable, "-c", _NO_PYSPARK_PROBE] - + [ipython_mode, pyspark_mode, tmp], - capture_output=True, - text=True, - env=env, + result = _run_python( + _NO_PYSPARK_PROBE, ipython_mode, pyspark_mode, tmp, env=env ) self.assertEqual(result.returncode, 0, result.stderr) self.assertNotIn("Traceback", result.stderr) diff --git a/tests/unit/test_instrumentation.py b/tests/unit/test_instrumentation.py index a86ec52a..4f7e5580 100644 --- a/tests/unit/test_instrumentation.py +++ b/tests/unit/test_instrumentation.py @@ -786,6 +786,11 @@ def to_pandas(self, plan, observations): table, *_ = self._execute_and_fetch(plan, observations) return "pdf", "ei" + def _set_relation_in_plan(self, plan, relation): + # Newer pyspark's LogicalPlan.to_proto() delegates this to the client + # (which may compress the plan); older pyspark sets plan.root itself. + plan.root.CopyFrom(relation) + def test_client_to_pandas_wrappers(inst): client = FakeClient() diff --git a/tests/unit/test_proxy.py b/tests/unit/test_proxy.py index ac778ea5..7235c5a3 100644 --- a/tests/unit/test_proxy.py +++ b/tests/unit/test_proxy.py @@ -40,7 +40,17 @@ def outgoing_socket_timeout(request): return request.param -@pytest.fixture(params=[0, 0.1, 0.2]) +# Nonzero waits make each test take seconds, while the zero waits still +# exercise every combination of socket timeouts, so only the zero waits are in +# the smoke tests. +_WAITS = [ + 0, + pytest.param(0.1, marks=pytest.mark.slow), + pytest.param(0.2, marks=pytest.mark.slow), +] + + +@pytest.fixture(params=_WAITS) def backend_wait(request): return request.param @@ -50,7 +60,7 @@ def incoming_socket_timeout(request): return request.param -@pytest.fixture(params=[0, 0.1, 0.2]) +@pytest.fixture(params=_WAITS) def client_wait(request): return request.param diff --git a/tests/unit/test_pypi_artifacts.py b/tests/unit/test_pypi_artifacts.py index 8c16a5d9..580d6646 100644 --- a/tests/unit/test_pypi_artifacts.py +++ b/tests/unit/test_pypi_artifacts.py @@ -10,6 +10,7 @@ import zipfile from packaging.requirements import InvalidRequirement +import pytest from google.cloud.managed_spark_connect.pypi_artifacts import PyPiArtifacts @@ -80,6 +81,7 @@ def _find_python_with_setuptools() -> str | None: return candidate return None + @pytest.mark.slow def test_sdist_and_wheel_exclude_oci_whiteouts_and_stale_egg_info(self): python_exe = self._find_python_with_setuptools() if python_exe is None: diff --git a/tests/unit/test_runtimes.py b/tests/unit/test_runtimes.py new file mode 100644 index 00000000..31fa81e5 --- /dev/null +++ b/tests/unit/test_runtimes.py @@ -0,0 +1,81 @@ +# 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 unittest + +from google.cloud.managed_spark_connect import _runtimes +from google.cloud.managed_spark_connect.session import ManagedSparkSession + + +class TestMajorMinor(unittest.TestCase): + + def test_valid_versions(self): + for version_str, expected in [ + ("3.0", (3, 0)), + ("3.0.12", (3, 0)), + ("4.0.1", (4, 0)), + ("4.1.0.dev1", (4, 1)), + ("4.2.0rc1", (4, 2)), + ]: + with self.subTest(version_str=version_str): + self.assertEqual(_runtimes.major_minor(version_str), expected) + + def test_invalid_versions(self): + for version_str in ["", "3", "unknown", "invalid.version"]: + with self.subTest(version_str=version_str): + self.assertIsNone(_runtimes.major_minor(version_str)) + + +class TestGetRuntimeInfo(unittest.TestCase): + + def test_known_runtime(self): + self.assertEqual( + _runtimes.get_runtime_info("3.0"), + _runtimes.RuntimeInfo(spark=(4, 0), python=(3, 12)), + ) + self.assertEqual( + _runtimes.get_runtime_info("3.0.12"), + _runtimes.get_runtime_info("3.0"), + ) + + def test_unknown_runtime(self): + for runtime_version in ["2.3", "99.0", "", "unknown"]: + with self.subTest(runtime_version=runtime_version): + self.assertIsNone(_runtimes.get_runtime_info(runtime_version)) + + def test_default_runtime_is_known(self): + self.assertIsNotNone( + _runtimes.get_runtime_info( + ManagedSparkSession._DEFAULT_RUNTIME_VERSION + ) + ) + + def test_default_runtime_matches_spark_extras(self): + """setup.py's [client] and [spark] extras pin the default runtime's Spark.""" + import pathlib + import re + + setup_py = (pathlib.Path(__file__).parents[2] / "setup.py").read_text( + encoding="utf-8" + ) + spark = _runtimes.format_version( + _runtimes.get_runtime_info( + ManagedSparkSession._DEFAULT_RUNTIME_VERSION + ).spark + ) + for requirement in ["pyspark-client", r"pyspark\[connect\]"]: + with self.subTest(requirement=requirement): + self.assertRegex( + setup_py, + re.compile(rf'"{requirement}~={re.escape(spark)}\.0"'), + ) diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index 7d64d779..d4a2d2bc 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -47,6 +47,7 @@ from pyspark.sql.connect.client.core import ConfigResult from pyspark.sql.connect.proto import Command, ConfigResponse, ExecutePlanRequest, Plan, Relation, SQL, SqlCommand, UserContext +import pytest from unittest import mock _MANAGED_SPARK_SESSIONS_BASE_URL = ( @@ -3040,6 +3041,7 @@ def test_wait_for_session_available_success( self.assertEqual(mock_client.get_session.call_count, 3) self.assertEqual(mock_sleep.call_count, 2) + @pytest.mark.slow @mock.patch("time.sleep", return_value=None) @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") def test_wait_for_session_available_timeout( @@ -3558,6 +3560,7 @@ def test_session_terminating_waits_then_deletes_and_warns( warning_msg = mock_logger.warning.call_args[0][0] self.assertIn("TERMINATING", warning_msg) + @pytest.mark.slow @mock.patch("google.auth.default") @mock.patch("google.cloud.dataproc_v1.SessionControllerClient") @mock.patch(