refactor(rust): colocate native wheel contract checks

.github/workflows/test-rust.yml

@@ -7,8 +7,8 @@ on:
       - ".cargo/**"
       - "pyproject.toml"
       - "rust-toolchain.toml"
-      - ".github/scripts/smoke_test_native_wheel.py"
-      - ".github/scripts/verify_linux_native_wheel.py"
+      - "litellm/rust_bridge/smoke_test_native_wheel.py"
+      - "litellm/rust_bridge/verify_linux_native_wheel.py"
       - ".github/workflows/test-rust.yml"
   pull_request:
     branches:
@@ -21,8 +21,8 @@ on:
       - ".cargo/**"
       - "pyproject.toml"
       - "rust-toolchain.toml"
-      - ".github/scripts/smoke_test_native_wheel.py"
-      - ".github/scripts/verify_linux_native_wheel.py"
+      - "litellm/rust_bridge/smoke_test_native_wheel.py"
+      - "litellm/rust_bridge/verify_linux_native_wheel.py"
       - ".github/workflows/test-rust.yml"
 
 permissions:
@@ -115,9 +115,9 @@ jobs:
           --config-setting "maturin.build-args=--features panic-test,extension-module"
 
       - name: Smoke-test native panic unwinding
-        run: python .github/scripts/smoke_test_native_wheel.py panic-dist/*.whl
+        run: python litellm/rust_bridge/smoke_test_native_wheel.py panic-dist/*.whl
 
       - name: Verify stripped native extension
         env:
           RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }}
-        run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl
+        run: python litellm/rust_bridge/verify_linux_native_wheel.py dist/*.whl

litellm/rust_bridge/verify_linux_native_wheel.py

@@ -6,18 +6,40 @@ import re
 import subprocess
 import sys
 import zipfile
+from collections.abc import Callable, Mapping, Sequence
 from email import policy
 from email.parser import BytesParser
 from itertools import product
 from pathlib import Path, PurePosixPath
 from types import ModuleType
-from typing import Final, cast
+from typing import Final, Protocol, cast
 
 EXPECTED_PYTHON_TAG: Final = "cp310"
 EXPECTED_ABI_TAG: Final = "abi3"
 EXPECTED_PLATFORM_TAG: Final = "linux_x86_64"
 
 
+class CommandRunner(Protocol):
+    def __call__(
+        self,
+        command: tuple[str, ...],
+        *,
+        check: bool,
+        capture_output: bool,
+        text: bool,
+    ) -> subprocess.CompletedProcess[str]: ...
+
+
+def _run_command(
+    command: tuple[str, ...],
+    *,
+    check: bool,
+    capture_output: bool,
+    text: bool,
+) -> subprocess.CompletedProcess[str]:
+    return subprocess.run(command, check=check, capture_output=capture_output, text=text)
+
+
 def _dist_info_directory(member: zipfile.ZipInfo) -> str | None:
     parts: Final = PurePosixPath(member.filename).parts
     if not parts or not parts[0].endswith(".dist-info"):
@@ -46,12 +68,19 @@ def _load_native_module(native_path: Path) -> ModuleType | None:
     return native_module
 
 
-def main() -> int:
-    if len(sys.argv) != 2:
-        sys.stderr.write(f"usage: {Path(sys.argv[0]).name} WHEEL\n")
+def main(
+    argv: Sequence[str] | None = None,
+    environment: Mapping[str, str] | None = None,
+    load_native_module: Callable[[Path], ModuleType | None] = _load_native_module,
+    run_command: CommandRunner = _run_command,
+) -> int:
+    arguments: Final = tuple(sys.argv if argv is None else argv)
+    resolved_environment: Final = os.environ if environment is None else environment
+    if len(arguments) != 2:
+        sys.stderr.write(f"usage: {Path(arguments[0]).name} WHEEL\n")
         return 2
 
-    wheel: Final = Path(sys.argv[1])
+    wheel: Final = Path(arguments[1])
     wheel_tags: Final = wheel.stem.rsplit("-", maxsplit=3)
     if len(wheel_tags) != 4:
         sys.stderr.write(f"cannot parse wheel tags from {wheel.name}\n")
@@ -110,8 +139,10 @@ def main() -> int:
         len(wheel_metadata_tags) == len(expanded_filename_tags)
         and frozenset(wheel_metadata_tags) == expanded_filename_tags
     )
-    commit_sha: Final = os.environ.get("RELEASE_WHEEL_COMMIT_SHA", os.environ.get("GITHUB_SHA", "unknown"))
-    rustc_version: Final = subprocess.run(
+    commit_sha: Final = resolved_environment.get(
+        "RELEASE_WHEEL_COMMIT_SHA", resolved_environment.get("GITHUB_SHA", "unknown")
+    )
+    rustc_version: Final = run_command(
         ("rustc", "--version"),
         check=True,
         capture_output=True,
@@ -147,14 +178,14 @@ def main() -> int:
             "",
         )
     )
-    summary_path: Final = os.environ.get("GITHUB_STEP_SUMMARY")
+    summary_path: Final = resolved_environment.get("GITHUB_STEP_SUMMARY")
     if summary_path is None:
         sys.stdout.write(size_report)
     else:
         Path(summary_path).write_text(size_report)
 
-    sections: Final = subprocess.run(
-        ("readelf", "--sections", "--wide", native_path),
+    sections: Final = run_command(
+        ("readelf", "--sections", "--wide", str(native_path)),
         check=True,
         capture_output=True,
         text=True,
@@ -163,14 +194,14 @@ def main() -> int:
     debug_sections_absent: Final = not debug_sections
     static_symbol_table_absent: Final = ".symtab" not in sections
 
-    dynamic_symbols: Final = subprocess.run(
-        ("readelf", "--dyn-syms", "--wide", native_path),
+    dynamic_symbols: Final = run_command(
+        ("readelf", "--dyn-syms", "--wide", str(native_path)),
         check=True,
         capture_output=True,
         text=True,
     ).stdout
     extension_entry_point_present: Final = "PyInit__native" in dynamic_symbols
-    native_module: Final = _load_native_module(native_path)
+    native_module: Final = load_native_module(native_path)
     native_module_loads: Final = native_module is not None
     panic_test_hook_absent: Final = native_module is not None and not hasattr(native_module, "_panic_for_test")
     native_size_limit: Final = 20_000_000

tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py

@@ -1,32 +1,16 @@
 from __future__ import annotations
 
-import importlib.util
 import subprocess
-import sys
 import zipfile
-from collections.abc import Callable
 from pathlib import Path
 from types import ModuleType
-from typing import Final, Protocol, cast
+from typing import Final
 
 import pytest
 
-_REPO_ROOT: Final = Path(__file__).resolve().parents[2]
-_MODULE_PATH: Final = _REPO_ROOT / ".github" / "scripts" / "verify_linux_native_wheel.py"
+from litellm.rust_bridge import verify_linux_native_wheel as verifier
 
-
-class _VerifierModule(Protocol):
-    subprocess: ModuleType
-    _load_native_module: Callable[[Path], ModuleType | None]
-    main: Callable[[], int]
-
-
-_SPEC: Final = importlib.util.spec_from_file_location("verify_linux_native_wheel", _MODULE_PATH)
-assert _SPEC is not None and _SPEC.loader is not None
-_LOADED_VERIFIER: Final = importlib.util.module_from_spec(_SPEC)
-sys.modules[_SPEC.name] = _LOADED_VERIFIER
-_SPEC.loader.exec_module(_LOADED_VERIFIER)
-verifier: Final = cast(_VerifierModule, _LOADED_VERIFIER)
+_MODULE_PATH: Final = Path(verifier.__file__)
 
 _EXPECTED_TAG: Final = "cp310-abi3-linux_x86_64"
 _NATIVE_MEMBER: Final = "litellm/rust_bridge/_native.abi3.so"
@@ -76,7 +60,6 @@ def _fake_subprocess_run(command: tuple[str, ...], **_: object) -> subprocess.Co
 
 
 def _run_verifier(
-    monkeypatch: pytest.MonkeyPatch,
     wheel: Path,
     *,
     exposes_panic: bool = False,
@@ -88,31 +71,33 @@ def _run_verifier(
     def _fake_load_native_module(_: Path) -> ModuleType:
         return native_module
 
-    monkeypatch.setattr(verifier, "_load_native_module", _fake_load_native_module)
-    monkeypatch.setattr(verifier.subprocess, "run", _fake_subprocess_run)
-    monkeypatch.setattr(sys, "argv", [str(_MODULE_PATH), str(wheel)])
-    monkeypatch.setenv("GITHUB_STEP_SUMMARY", str(wheel.parent / "summary.md"))
-    return verifier.main()
+    environment: Final = {"GITHUB_STEP_SUMMARY": str(wheel.parent / "summary.md")}
+    return verifier.main(
+        (str(_MODULE_PATH), str(wheel)),
+        environment,
+        _fake_load_native_module,
+        _fake_subprocess_run,
+    )
 
 
-def test_accepts_expected_release_wheel_tags(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+def test_accepts_expected_release_wheel_tags(tmp_path: Path) -> None:
     wheel: Final = _write_wheel(tmp_path, filename_tag=_EXPECTED_TAG)
 
-    assert _run_verifier(monkeypatch, wheel) == 0
+    assert _run_verifier(wheel) == 0
 
 
-def test_rejects_cp312_version_specific_wheel(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+def test_rejects_cp312_version_specific_wheel(tmp_path: Path) -> None:
     tag: Final = "cp312-cp312-linux_x86_64"
     wheel: Final = _write_wheel(tmp_path, filename_tag=tag, metadata_tags=(tag,))
 
-    assert _run_verifier(monkeypatch, wheel) == 1
+    assert _run_verifier(wheel) == 1
 
 
-def test_rejects_non_linux_platform_tag(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+def test_rejects_non_linux_platform_tag(tmp_path: Path) -> None:
     tag: Final = "cp310-abi3-win_amd64"
     wheel: Final = _write_wheel(tmp_path, filename_tag=tag, metadata_tags=(tag,))
 
-    assert _run_verifier(monkeypatch, wheel) == 1
+    assert _run_verifier(wheel) == 1
 
 
 @pytest.mark.parametrize(
@@ -122,17 +107,15 @@ def test_rejects_non_linux_platform_tag(tmp_path: Path, monkeypatch: pytest.Monk
 )
 def test_rejects_missing_or_mismatched_wheel_metadata_tag(
     tmp_path: Path,
-    monkeypatch: pytest.MonkeyPatch,
     metadata_tags: tuple[str, ...] | None,
 ) -> None:
     wheel: Final = _write_wheel(tmp_path, filename_tag=_EXPECTED_TAG, metadata_tags=metadata_tags)
 
-    assert _run_verifier(monkeypatch, wheel) == 1
+    assert _run_verifier(wheel) == 1
 
 
 def test_rejects_wheel_metadata_from_wrong_dist_info_directory(
     tmp_path: Path,
-    monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     wheel: Final = _write_wheel(
         tmp_path,
@@ -140,20 +123,20 @@ def test_rejects_wheel_metadata_from_wrong_dist_info_directory(
         dist_info="decoy-1.0.0.dist-info",
     )
 
-    assert _run_verifier(monkeypatch, wheel) == 1
+    assert _run_verifier(wheel) == 1
 
 
-def test_rejects_duplicate_wheel_metadata_tags(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+def test_rejects_duplicate_wheel_metadata_tags(tmp_path: Path) -> None:
     wheel: Final = _write_wheel(
         tmp_path,
         filename_tag=_EXPECTED_TAG,
         metadata_tags=(_EXPECTED_TAG, _EXPECTED_TAG),
     )
 
-    assert _run_verifier(monkeypatch, wheel) == 1
+    assert _run_verifier(wheel) == 1
 
 
-def test_rejects_duplicate_wheel_metadata_file(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+def test_rejects_duplicate_wheel_metadata_file(tmp_path: Path) -> None:
     with pytest.warns(UserWarning, match="Duplicate name"):
         wheel: Final = _write_wheel(
             tmp_path,
@@ -161,10 +144,10 @@ def test_rejects_duplicate_wheel_metadata_file(tmp_path: Path, monkeypatch: pyte
             duplicate_wheel=True,
         )
 
-    assert _run_verifier(monkeypatch, wheel) == 1
+    assert _run_verifier(wheel) == 1
 
 
-def test_rejects_production_module_exposing_panic_hook(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+def test_rejects_production_module_exposing_panic_hook(tmp_path: Path) -> None:
     wheel: Final = _write_wheel(tmp_path, filename_tag=_EXPECTED_TAG)
 
-    assert _run_verifier(monkeypatch, wheel, exposes_panic=True) == 1
+    assert _run_verifier(wheel, exposes_panic=True) == 1