test: handle pending kernel geospatial release

Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com>

Author: cathleeny

src/databricks/sql/backend/kernel/client.py

@@ -179,8 +179,7 @@ def _kernel_geospatial_kwargs(value: bool) -> Dict[str, bool]:
     """
     if not isinstance(value, bool):
         raise ValueError(
-            "enable_geospatial_support must be a bool; "
-            f"got {type(value).__name__}"
+            "enable_geospatial_support must be a bool; " f"got {type(value).__name__}"
         )
     if not _kernel_session_accepts_kwarg("enable_geospatial_support"):
         raise NotSupportedError(
@@ -288,9 +287,7 @@ def __init__(
         # bytes}`` through pyarrow; False requests WKT / EWKT strings. This is
         # intentionally separate from ``session_configuration``: it is never
         # forwarded to SEA.
-        self._enable_geospatial_support = kwargs.get(
-            "enable_geospatial_support", True
-        )
+        self._enable_geospatial_support = kwargs.get("enable_geospatial_support", True)
         # Kernel telemetry phase 7 adds binding/runtime identity and
         # telemetry config kwargs directly to ``databricks_sql_kernel.Session``.
         self._telemetry_options = kwargs.get("telemetry_options") or {}

src/databricks/sql/session.py

@@ -314,9 +314,7 @@ def _create_backend(
                 retry_options=kernel_retry_options,
                 request_timeout_secs=kwargs.get("_socket_timeout"),
                 max_connections=kwargs.get("_pool_maxsize") or None,
-                enable_geospatial_support=kwargs.get(
-                    "enable_geospatial_support", True
-                ),
+                enable_geospatial_support=kwargs.get("enable_geospatial_support", True),
                 telemetry_options=kernel_telemetry_options,
             )
 

tests/unit/test_session.py

@@ -796,14 +796,23 @@ def test_connect_use_kernel_instantiates_real_kernel_backend(self):
     def test_enable_geospatial_support_matches_real_kernel_signature(self):
         self._real_kernel_or_skip()
 
-        from databricks.sql.backend.kernel.client import _kernel_geospatial_kwargs
-
-        assert _kernel_geospatial_kwargs(True) == {
-            "enable_geospatial_support": True
-        }
-        assert _kernel_geospatial_kwargs(False) == {
-            "enable_geospatial_support": False
-        }
+        from databricks.sql.backend.kernel.client import (
+            _kernel_geospatial_kwargs,
+            _kernel_session_accepts_kwarg,
+        )
+
+        # The ordinary kernel unit-test tier installs the latest published
+        # wheel, which may lag the KERNEL_REV source pin while the matching
+        # kernel change is still in flight. Validate the compatibility error
+        # in that case; once the wheel includes the typed option, validate
+        # both values against its real PyO3 signature.
+        if not _kernel_session_accepts_kwarg("enable_geospatial_support"):
+            with pytest.raises(NotSupportedError, match="newer databricks-sql-kernel"):
+                _kernel_geospatial_kwargs(True)
+            return
+
+        assert _kernel_geospatial_kwargs(True) == {"enable_geospatial_support": True}
+        assert _kernel_geospatial_kwargs(False) == {"enable_geospatial_support": False}
 
 
 class TestReydenThriftFallback: