import types import pytest from producer import lazy PYPI_JSON = { "urls": [ { "filename": "torch-2.7.1-cp311-cp311-manylinux_2_28_x86_64.whl", "url": "https://files.pythonhosted.org/packages/xx/torch-2.7.1-cp311-cp311-manylinux_2_28_x86_64.whl", "digests": {"sha256": "abc123"}, "size": 766_668_798, }, { "filename": "torch-2.7.1-cp311-cp311-win_amd64.whl", "url": "https://example.com/win.whl", "digests": {}, "size": 1, }, { "filename": "torch-2.7.1.tar.gz", "url": "https://example.com/src.tar.gz", "digests": {}, "size": 1, }, ] } _SHA = "e1a846516570851234567890abcdef1234567890abcdef1234567890abcdef12" GPU_HTML = ( "" 'old new' 'win' 'relative-dup' "" ) def test_choose_pypi_picks_linux_wheel(): got = lazy._choose_pypi(PYPI_JSON) assert got is not None name, url, sha, size = got assert name.endswith("manylinux_2_28_x86_64.whl") assert "pythonhosted" in url assert sha == "abc123" assert size == 766_668_798 def test_choose_gpu_matches_version_and_arch(): got = lazy._choose_gpu(GPU_HTML, "torch", "2.7.1") assert got is not None name, url, sha, size = got assert name == "torch-2.7.1+cu126-cp311-cp311-manylinux_2_28_x86_64.whl" assert url == ( "https://download-r2.pytorch.org/whl/cu126/" "torch-2.7.1%2Bcu126-cp311-cp311-manylinux_2_28_x86_64.whl" ) assert sha == _SHA assert size is None def test_choose_gpu_no_match_returns_none(): assert lazy._choose_gpu("", "torch", "2.7.1") is None def test_resolve_torch_wheels_needs_both_packages(monkeypatch): def fake(pkg, ver): return None if pkg == "torchaudio" else ("t.whl", "u", None, 1) monkeypatch.setattr(lazy, "_pypi_wheel", fake) assert lazy._resolve_torch_wheels(gpu=False) == [] def test_resolve_torch_wheels_returns_both(monkeypatch): wheels = [("torch.whl", "u1", "s1", 1), ("torchaudio.whl", "u2", None, None)] monkeypatch.setattr(lazy, "_pypi_wheel", lambda pkg, ver: wheels.pop(0)) got = lazy._resolve_torch_wheels(gpu=False) assert [w[0] for w in got] == ["torch.whl", "torchaudio.whl"] def test_ensure_torch_downloads_wheels_then_installs(monkeypatch, tmp_path): wheels = [ ("torch-2.7.1-cp311.whl", "https://x/torch.whl", "sha", 10), ("torchaudio-2.7.1-cp311.whl", "https://x/ta.whl", None, 5), ] monkeypatch.setattr(lazy, "has_module", lambda name: False) monkeypatch.setattr(lazy, "_resolve_torch_wheels", lambda gpu: wheels) monkeypatch.setattr(lazy, "find_uv", lambda: "uv") monkeypatch.setattr(lazy, "WHEEL_CACHE", tmp_path / "wheels") downloads: list[str] = [] runs: list[tuple[list[str], str]] = [] def fake_download(url, dest, label=None, expected_size=None, sha256=None, timeout=60.0): downloads.append(dest.name) return dest monkeypatch.setattr(lazy.ui, "download", fake_download) monkeypatch.setattr(lazy.ui, "run", lambda cmd, label, check=True: runs.append((cmd, label))) lazy.ensure_torch() assert downloads == ["torch-2.7.1-cp311.whl", "torchaudio-2.7.1-cp311.whl"] assert len(runs) == 1 cmd = runs[0][0] assert cmd[0] == "uv" and cmd[1] == "pip" and cmd[2] == "install" assert str(tmp_path / "wheels" / "torch-2.7.1-cp311.whl") in cmd assert str(tmp_path / "wheels" / "torchaudio-2.7.1-cp311.whl") in cmd def test_ensure_torch_falls_back_to_uv_index(monkeypatch): monkeypatch.setattr(lazy, "has_module", lambda name: False) monkeypatch.setattr(lazy, "_resolve_torch_wheels", lambda gpu: []) monkeypatch.setattr(lazy, "find_uv", lambda: "uv") monkeypatch.setattr(lazy, "gpu_present", lambda: True) runs: list[list[str]] = [] monkeypatch.setattr(lazy.ui, "run", lambda cmd, label, check=True: runs.append(cmd)) lazy.ensure_torch() cmd = runs[0] assert "torch==2.7.1+cu126" in cmd assert "--index-url" in cmd assert lazy.TORCH_GPU_INDEX in cmd def test_ensure_torch_falls_back_to_pypi_cpu(monkeypatch): monkeypatch.setattr(lazy, "has_module", lambda name: False) monkeypatch.setattr(lazy, "_resolve_torch_wheels", lambda gpu: []) monkeypatch.setattr(lazy, "find_uv", lambda: "uv") monkeypatch.setattr(lazy, "gpu_present", lambda: False) runs: list[list[str]] = [] monkeypatch.setattr(lazy.ui, "run", lambda cmd, label, check=True: runs.append(cmd)) lazy.ensure_torch() cmd = runs[0] assert "torch==2.7.1" in cmd assert "--index-url" not in cmd def test_ensure_torch_noop_when_installed(monkeypatch): monkeypatch.setattr(lazy, "has_module", lambda name: True) runs: list = [] monkeypatch.setattr(lazy.ui, "run", lambda cmd, label, check=True: runs.append(cmd)) lazy.ensure_torch() assert runs == [] def test_ensure_import_noop_when_importable(monkeypatch): installs: list[list[str]] = [] monkeypatch.setattr(lazy, "ensure", lambda pkgs, **_k: installs.append(list(pkgs))) lazy.ensure_import("numpy", purpose="test") assert installs == [] def test_ensure_import_installs_missing_module(monkeypatch): installs: list[list[str]] = [] monkeypatch.setattr(lazy, "ensure", lambda pkgs, **_k: installs.append(list(pkgs))) state = {"round": 0} def fake_import(name): if name == "fakeengine": if state["round"] == 0: state["round"] = 1 raise ModuleNotFoundError("No module named 'addict'", name="addict") return types.ModuleType("fakeengine") raise ModuleNotFoundError(f"No module named {name!r}", name=name) monkeypatch.setattr(lazy.importlib, "import_module", fake_import) lazy.ensure_import("fakeengine", purpose="test") assert installs == [["addict"]] def test_ensure_import_dotted_missing_reraises_without_install(monkeypatch): installs: list[list[str]] = [] monkeypatch.setattr(lazy, "ensure", lambda pkgs, **_k: installs.append(list(pkgs))) def fake_import(name): raise ModuleNotFoundError("No module named 'pkg.sub'", name="pkg.sub") monkeypatch.setattr(lazy.importlib, "import_module", fake_import) with pytest.raises(ModuleNotFoundError): lazy.ensure_import("whatever", purpose="test") assert installs == [] def test_ensure_import_gives_up_after_rounds(monkeypatch): installs: list[list[str]] = [] monkeypatch.setattr(lazy, "ensure", lambda pkgs, **_k: installs.append(list(pkgs))) def fake_import(name): raise ModuleNotFoundError("No module named 'ghost'", name="ghost") monkeypatch.setattr(lazy.importlib, "import_module", fake_import) with pytest.raises(lazy.EngineUnavailable, match="ghost"): lazy.ensure_import("whatever", purpose="test") assert len(installs) == 4 assert all(pkgs == ["ghost"] for pkgs in installs) def test_ensure_import_maps_pil_to_pillow(monkeypatch): installs: list[list[str]] = [] monkeypatch.setattr(lazy, "ensure", lambda pkgs, **_k: installs.append(list(pkgs))) state = {"round": 0} def fake_import(name): if name == "fakeengine": if state["round"] == 0: state["round"] = 1 raise ModuleNotFoundError("No module named 'PIL'", name="PIL") return types.ModuleType("fakeengine") raise ModuleNotFoundError(f"No module named {name!r}", name=name) monkeypatch.setattr(lazy.importlib, "import_module", fake_import) lazy.ensure_import("fakeengine", purpose="test") assert installs == [["pillow"]] def test_ensure_call_passthrough_without_install(monkeypatch): installs: list[list[str]] = [] monkeypatch.setattr(lazy, "ensure", lambda pkgs, **_k: installs.append(list(pkgs))) assert lazy.ensure_call(lambda: 42, purpose="test") == 42 assert installs == [] def test_ensure_call_installs_missing_module(monkeypatch): installs: list[list[str]] = [] monkeypatch.setattr(lazy, "ensure", lambda pkgs, **_k: installs.append(list(pkgs))) state = {"round": 0} def fake_fn(): if state["round"] == 0: state["round"] = 1 raise ModuleNotFoundError("No module named 'addict'", name="addict") return "ok" assert lazy.ensure_call(fake_fn, purpose="test") == "ok" assert installs == [["addict"]] def test_ensure_call_dotted_missing_reraises_without_install(monkeypatch): installs: list[list[str]] = [] monkeypatch.setattr(lazy, "ensure", lambda pkgs, **_k: installs.append(list(pkgs))) def fake_fn(): raise ModuleNotFoundError("No module named 'pkg.sub'", name="pkg.sub") with pytest.raises(ModuleNotFoundError): lazy.ensure_call(fake_fn, purpose="test") assert installs == [] def test_ensure_call_unnamed_missing_reraises_without_install(monkeypatch): installs: list[list[str]] = [] monkeypatch.setattr(lazy, "ensure", lambda pkgs, **_k: installs.append(list(pkgs))) def fake_fn(): raise ModuleNotFoundError("boom") with pytest.raises(ModuleNotFoundError): lazy.ensure_call(fake_fn, purpose="test") assert installs == [] def test_ensure_call_maps_pil_to_pillow(monkeypatch): installs: list[list[str]] = [] monkeypatch.setattr(lazy, "ensure", lambda pkgs, **_k: installs.append(list(pkgs))) state = {"round": 0} def fake_fn(): if state["round"] == 0: state["round"] = 1 raise ModuleNotFoundError("No module named 'PIL'", name="PIL") return "ok" assert lazy.ensure_call(fake_fn, purpose="test") == "ok" assert installs == [["pillow"]] def test_ensure_call_gives_up_after_rounds(monkeypatch): installs: list[list[str]] = [] monkeypatch.setattr(lazy, "ensure", lambda pkgs, **_k: installs.append(list(pkgs))) def fake_fn(): raise ModuleNotFoundError("No module named 'ghost'", name="ghost") with pytest.raises(lazy.EngineUnavailable, match="ghost"): lazy.ensure_call(fake_fn, purpose="test") assert len(installs) == 8 assert all(pkgs == ["ghost"] for pkgs in installs)