test: annotate the shared test fixtures and pool tests

tests/_ray_test_support.py

@@ -7,6 +7,7 @@
 from __future__ import annotations
 
 import threading
+from typing import Any, Optional
 
 
 def patch_memory_profiler() -> None:
@@ -24,7 +25,7 @@ def patch_memory_profiler() -> None:
 
     original_init = ray_data_util.MemoryProfiler.__init__
 
-    def safe_init(self, poll_interval_s):
+    def safe_init(self: Any, poll_interval_s: Optional[float]) -> None:
         self._poll_interval_s = poll_interval_s
         try:
             original_init(self, poll_interval_s)
@@ -35,10 +36,15 @@ def safe_init(self, poll_interval_s):
             self._uss_poll_thread = None
             self._stop_uss_poll_event = None
 
-    ray_data_util.MemoryProfiler.__init__ = safe_init
-
+    # Monkey-patching methods is the point of this module, so the
+    # ``method-assign`` guard does not apply here.
+    ray_data_util.MemoryProfiler.__init__ = safe_init  # type: ignore[method-assign]
     if hasattr(ray_data_util.MemoryProfiler, "_can_estimate_uss"):
-        ray_data_util.MemoryProfiler._can_estimate_uss = lambda self: False
+
+        def _cannot_estimate_uss(self: Any) -> bool:
+            return False
+
+        ray_data_util.MemoryProfiler._can_estimate_uss = _cannot_estimate_uss  # type: ignore[method-assign,assignment]
 
 
 def patch_psutil_for_containers() -> None:
@@ -62,7 +68,7 @@ def patch_psutil_for_containers() -> None:
 
     original_fn = uv_runtime_env_hook._get_uv_run_cmdline
 
-    def _safe_get_uv_run_cmdline():
+    def _safe_get_uv_run_cmdline() -> Any:
         try:
             return original_fn()
         except Exception:
@@ -91,20 +97,20 @@ def _patch_worker_log_offset() -> None:
     _orig_out = Worker.get_current_out_offset
     _orig_err = Worker.get_current_err_offset
 
-    def _safe_out_offset(self) -> int:
+    def _safe_out_offset(self: Any) -> int:
         try:
             return _orig_out(self)
         except FileNotFoundError:
             return 0
 
-    def _safe_err_offset(self) -> int:
+    def _safe_err_offset(self: Any) -> int:
         try:
             return _orig_err(self)
         except FileNotFoundError:
             return 0
 
-    Worker.get_current_out_offset = _safe_out_offset
-    Worker.get_current_err_offset = _safe_err_offset
+    Worker.get_current_out_offset = _safe_out_offset  # type: ignore[method-assign]
+    Worker.get_current_err_offset = _safe_err_offset  # type: ignore[method-assign]
 
 
 def setup_worker() -> None:

tests/conftest.py

@@ -4,6 +4,7 @@
 import shutil
 import sys
 import tempfile
+from collections.abc import Iterator
 
 import pytest
 import ray
@@ -21,7 +22,7 @@
 
 
 @pytest.fixture(scope="session", autouse=True)
-def ray_context():
+def ray_context() -> Iterator[None]:
     """Initialize Ray once per pytest session.
 
     Defined in conftest.py so that running multiple test files concurrently

tests/test_field_path.py

@@ -7,7 +7,7 @@
 )
 
 
-def _nested_schema():
+def _nested_schema() -> pa.Schema:
     return pa.schema(
         [
             pa.field("id", pa.int64()),
@@ -31,20 +31,20 @@ def _nested_schema():
     )
 
 
-def test_parse_field_path_supports_quoted_literal_dot_segments():
+def test_parse_field_path_supports_quoted_literal_dot_segments() -> None:
     assert parse_field_path("meta.`a.b`") == ["meta", "a.b"]
     assert parse_field_path("`meta`.`userId`") == ["meta", "userId"]
     assert parse_field_path("meta.`a``b`") == ["meta", "a`b"]
 
 
-def test_canonical_field_path_matches_lance_formatting_rules():
+def test_canonical_field_path_matches_lance_formatting_rules() -> None:
     assert canonical_field_path("`meta`.`userId`") == "meta.userId"
     assert canonical_field_path("meta.`a.b`") == "meta.`a.b`"
     assert canonical_field_path("`meta-data`.`user-id`") == "`meta-data`.`user-id`"
     assert canonical_field_path("meta.`a``b`") == "meta.`a``b`"
 
 
-def test_resolve_arrow_field_path_supports_nested_and_literal_dot():
+def test_resolve_arrow_field_path_supports_nested_and_literal_dot() -> None:
     schema = _nested_schema()
 
     assert resolve_arrow_field_path(schema, "meta.userId").field.name == "userId"
@@ -58,7 +58,7 @@ def test_resolve_arrow_field_path_supports_nested_and_literal_dot():
     assert resolved_hyphen.field.name == "user-id"
 
 
-def test_resolve_arrow_field_path_requires_disambiguated_same_leaf_name():
+def test_resolve_arrow_field_path_requires_disambiguated_same_leaf_name() -> None:
     schema = _nested_schema()
 
     with pytest.raises(KeyError):
@@ -68,7 +68,7 @@ def test_resolve_arrow_field_path_requires_disambiguated_same_leaf_name():
     assert resolve_arrow_field_path(schema, "other.leaf").field.name == "leaf"
 
 
-def test_parse_field_path_rejects_malformed_quoted_paths():
+def test_parse_field_path_rejects_malformed_quoted_paths() -> None:
     with pytest.raises(ValueError, match="unterminated"):
         parse_field_path("meta.`a.b")
     with pytest.raises(ValueError, match="empty path segment"):

tests/test_pool.py

@@ -1,19 +1,23 @@
 import logging
+from typing import Any, Optional
 
+import pytest
 from lance_ray import pool as pool_mod
 
 
-def test_init_global_pool_reuses_existing_pool(monkeypatch):
-    events = []
+def test_init_global_pool_reuses_existing_pool(monkeypatch: pytest.MonkeyPatch) -> None:
+    events: list[str | tuple[str, int, Optional[dict[str, Any]]]] = []
 
     class FakePool:
-        def __init__(self, processes, ray_remote_args):
+        def __init__(
+            self, processes: int, ray_remote_args: Optional[dict[str, Any]]
+        ) -> None:
             events.append(("init", processes, ray_remote_args))
 
-        def close(self):
+        def close(self) -> None:
             events.append("close")
 
-        def join(self):
+        def join(self) -> None:
             events.append("join")
 
     pool_mod.clear_global_pool()
@@ -39,11 +43,15 @@ def join(self):
     ]
 
 
-def test_init_global_pool_warns_when_existing_pool_size_differs(monkeypatch, caplog):
-    events = []
+def test_init_global_pool_warns_when_existing_pool_size_differs(
+    monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
+) -> None:
+    events: list[tuple[str, int, Optional[dict[str, Any]]]] = []
 
     class FakePool:
-        def __init__(self, processes, ray_remote_args):
+        def __init__(
+            self, processes: int, ray_remote_args: Optional[dict[str, Any]] = None
+        ) -> None:
             events.append(("init", processes, ray_remote_args))
 
     pool_mod.clear_global_pool()
@@ -59,7 +67,9 @@ def __init__(self, processes, ray_remote_args):
     pool_mod.clear_global_pool()
 
 
-def test_get_or_create_pool_warns_when_global_pool_size_differs(caplog):
+def test_get_or_create_pool_warns_when_global_pool_size_differs(
+    caplog: pytest.LogCaptureFixture,
+) -> None:
     class FakePool:
         processes = 4
 
@@ -77,14 +87,14 @@ class FakePool:
     assert "requested 16 workers will be ignored" in caplog.text
 
 
-def test_set_global_pool_can_clear_without_closing():
-    events = []
+def test_set_global_pool_can_clear_without_closing() -> None:
+    events: list[str] = []
 
     class FakePool:
-        def close(self):
+        def close(self) -> None:
             events.append("close")
 
-        def join(self):
+        def join(self) -> None:
             events.append("join")
 
     pool = FakePool()