SnapOtter/packages/ai/python/tests/test_install_feature_download.py

import importlib.util
import os
import sys
import types


def load_installer():
    script_path = os.path.join(os.path.dirname(__file__), "..", "install_feature.py")
    spec = importlib.util.spec_from_file_location("install_feature_under_test", script_path)
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    spec.loader.exec_module(module)
    return module


def test_main_temporarily_lifts_hugging_face_offline_flags(monkeypatch):
    """An explicit bundle install stays online while runtime remains fail-closed."""
    installer = load_installer()
    monkeypatch.setenv("HF_HUB_OFFLINE", "1")
    monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1")
    observed = {}

    def fake_install():
        observed["hf"] = os.environ["HF_HUB_OFFLINE"]
        observed["transformers"] = os.environ["TRANSFORMERS_OFFLINE"]

    monkeypatch.setattr(installer, "_install", fake_install)

    installer.main()

    assert observed == {"hf": "0", "transformers": "0"}
    assert os.environ["HF_HUB_OFFLINE"] == "1"
    assert os.environ["TRANSFORMERS_OFFLINE"] == "1"


def test_download_with_hf_hub_uses_accelerated_client(monkeypatch, tmp_path):
    installer = load_installer()
    downloaded = tmp_path / "hf-cache" / "bundle.tar.gz"
    downloaded.parent.mkdir()

    calls = {}

    def fake_hf_hub_download(**kwargs):
        calls.update(kwargs)
        downloaded.write_bytes(b"archive")
        return str(downloaded)

    fake_module = types.ModuleType("huggingface_hub")
    fake_module.hf_hub_download = fake_hf_hub_download
    monkeypatch.setitem(sys.modules, "huggingface_hub", fake_module)

    progress = []
    monkeypatch.setattr(installer, "emit_progress", lambda p, s: progress.append((p, s)))

    dest = tmp_path / "staging" / "object-eraser-colorize-amd64-gpu.tar.gz"
    dest.parent.mkdir()

    assert (
        installer.download_with_hf_hub(
            "snapotter/feature-bundles",
            "v2.0.0/object-eraser-colorize-amd64-gpu.tar.gz",
            str(dest),
            100,
            2,
            85,
        )
        is True
    )

    assert dest.read_bytes() == b"archive"
    assert calls["repo_id"] == "snapotter/feature-bundles"
    assert calls["repo_type"] == "model"
    assert calls["filename"] == "v2.0.0/object-eraser-colorize-amd64-gpu.tar.gz"
    assert any("accelerated" in stage.lower() for _, stage in progress)


def test_download_with_hf_hub_cleans_cache_when_download_raises(monkeypatch, tmp_path):
    """A failed accelerated download must not leak its .cache staging tree onto
    disk before the urllib fallback runs."""
    installer = load_installer()
    staging = tmp_path / "staging"
    staging.mkdir()
    # Simulate a partial hf cache tree left behind by a failed transfer.
    leaked_cache = staging / ".cache" / "huggingface" / "download"
    leaked_cache.mkdir(parents=True)
    (leaked_cache / "blob.incomplete").write_bytes(b"partial")
    leaked_nested = staging / "v2.0.0"
    leaked_nested.mkdir()

    def fake_hf_hub_download(**_kwargs):
        raise RuntimeError("xet CAS unreachable")

    fake_module = types.ModuleType("huggingface_hub")
    fake_module.hf_hub_download = fake_hf_hub_download
    monkeypatch.setitem(sys.modules, "huggingface_hub", fake_module)
    monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)

    dest = staging / "object-eraser-colorize-amd64-gpu.tar.gz"

    assert (
        installer.download_with_hf_hub(
            "deepsafe/feature-bundles",
            "v2.0.0/object-eraser-colorize-amd64-gpu.tar.gz",
            str(dest),
            100,
            2,
            85,
        )
        is False
    )
    # Both the .cache tree and the nested archive dir are reclaimed.
    assert not (staging / ".cache").exists()
    assert not (staging / "v2.0.0").exists()


def test_ensure_hf_hub_noops_when_client_already_importable(monkeypatch, tmp_path):
    installer = load_installer()
    fake_module = types.ModuleType("huggingface_hub")
    monkeypatch.setitem(sys.modules, "huggingface_hub", fake_module)

    ran = {"pip": False}
    monkeypatch.setattr(
        installer.subprocess, "run", lambda *a, **k: ran.__setitem__("pip", True)
    )
    monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)

    installer.ensure_hf_hub(str(tmp_path))
    assert ran["pip"] is False


def test_ensure_hf_hub_self_heals_missing_client(monkeypatch, tmp_path):
    """On a drifted venv where huggingface_hub is missing, ensure_hf_hub must
    pip-install it into that venv rather than let the caller fall back to the
    slow single-stream urllib downloader silently."""
    installer = load_installer()
    monkeypatch.delitem(sys.modules, "huggingface_hub", raising=False)

    # Make the huggingface_hub import fail deterministically so ensure_hf_hub
    # takes its self-heal branch.
    import builtins

    real_import = builtins.__import__

    def blocked_import(name, *a, **k):
        if name == "huggingface_hub":
            raise ImportError("No module named 'huggingface_hub'")
        return real_import(name, *a, **k)

    monkeypatch.setattr(builtins, "__import__", blocked_import)

    venv = tmp_path / "venv"
    (venv / "bin").mkdir(parents=True)
    (venv / "bin" / "python3").write_text("")

    pip_calls = []
    monkeypatch.setattr(
        installer.subprocess,
        "run",
        lambda cmd, **k: pip_calls.append(cmd) or types.SimpleNamespace(returncode=0),
    )
    monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)

    installer.ensure_hf_hub(str(venv))

    assert len(pip_calls) == 1
    cmd = pip_calls[0]
    assert cmd[0] == str(venv / "bin" / "python3")
    assert "install" in cmd
    spec = next(part for part in cmd if part.startswith("huggingface-hub["))
    assert "hf_xet" in spec
    assert "hf_transfer" in spec