diff --git a/tests/_ray_test_support.py b/tests/_ray_test_support.py index 75a5717b..8b736f36 100644 --- a/tests/_ray_test_support.py +++ b/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: diff --git a/tests/conftest.py b/tests/conftest.py index 3814fee7..9d3fb870 100644 --- a/tests/conftest.py +++ b/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 diff --git a/tests/test_field_path.py b/tests/test_field_path.py index 599b6221..d9316034 100644 --- a/tests/test_field_path.py +++ b/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"): diff --git a/tests/test_pool.py b/tests/test_pool.py index febfe53d..55728c18 100644 --- a/tests/test_pool.py +++ b/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()