Blame for tests/test_task_isolation.py

"""CPU-only regression tests for task isolation (no H3 checkpoint load)."""
from __future__ import annotations

import os
import pickle
import tempfile
import threading
import time
import uuid
from dataclasses import asdict
from pathlib import Path
from unittest import mock

import pytest

# Import only the lightweight pieces we need without triggering model load.
# app.py imports spaces/gradio at module level; patch minimally if needed.
import sys
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))


def test_task_id_unique_and_unpredictable():
    from app import GenerationTask
    ids = {GenerationTask.create().task_id for _ in range(50)}
    assert len(ids) == 50
    for tid in ids:
        assert len(tid) >= 16
        assert all(c in "0123456789abcdef" for c in tid)


def test_task_pickle_safe():
    from app import GenerationTask
    t = GenerationTask.create(prompt_len=10, canvas="960x544")
    t.applied_loras.append("turbo")
    data = pickle.dumps(t)
    t2 = pickle.loads(data)
    assert t2.task_id == t.task_id
    assert t2.applied_loras == ["turbo"]
    assert t2.work_dir == t.work_dir


def test_task_no_heavy_objects():
    from app import GenerationTask
    t = GenerationTask.create()
    d = asdict(t)
    for k, v in d.items():
        assert not hasattr(v, "cuda"), f"{k} looks like a tensor/device object"
        assert type(v).__name__ != "Lock"
        assert not isinstance(v, type(lambda: None))


def test_task_filesystem_isolation(tmp_path, monkeypatch):
    from app import GenerationTask, _TASK_TEMP_ROOT
    monkeypatch.setattr("app._TASK_TEMP_ROOT", str(tmp_path))
    a = GenerationTask.create()
    b = GenerationTask.create()
    assert a.work_dir != b.work_dir
    assert a.task_id in a.work_dir
    assert os.path.isdir(a.work_dir)
    assert os.path.isdir(os.path.join(a.work_dir, "intermediates"))
    # no collision
    open(os.path.join(a.work_dir, f"h3-{a.task_id}.mp4"), "w").write("a")
    open(os.path.join(b.work_dir, f"h3-{b.task_id}.mp4"), "w").write("b")
    assert open(os.path.join(a.work_dir, f"h3-{a.task_id}.mp4")).read() == "a"


def test_model_mutation_lock_is_reentrant():
    from app import _MODEL_MUTATION_LOCK
    with _MODEL_MUTATION_LOCK:
        with _MODEL_MUTATION_LOCK:
            pass  # reentrant


def test_export_rejects_outside_path(tmp_path, monkeypatch):
    from app import export_last_frame, _TASK_TEMP_ROOT
    monkeypatch.setattr("app._TASK_TEMP_ROOT", str(tmp_path))
    evil = "/etc/passwd"
    with pytest.raises(Exception):  # gr.Error or similar
        export_last_frame(evil)


def test_revert_token_not_in_argv(tmp_path, monkeypatch):
    """Regression: HF_TOKEN must not appear in printed commands or argv."""
    import revert_to_baseline as rtb
    sentinel = "hf_SENTINEL_TOKEN_DO_NOT_LEAK_xyz"
    monkeypatch.setenv("HF_TOKEN", sentinel)
    # Capture what run() would print / pass
    captured = []
    def fake_run(cmd, cwd=None, env=None):
        captured.append(list(cmd))
        # Do not actually run git
        return None
    monkeypatch.setattr(rtb, "run", fake_run)
    # Patch TemporaryDirectory to avoid real network
    class FakeTmp:
        def __enter__(self):
            d = str(tmp_path / "ask")
            os.makedirs(d, exist_ok=True)
            return d
        def __exit__(self, *a):
            pass
    monkeypatch.setattr(rtb.tempfile, "TemporaryDirectory", lambda: FakeTmp())
    # Also patch _write_askpass to not fail
    def fake_askpass(tmpdir):
        p = os.path.join(tmpdir, "askpass.sh")
        open(p, "w").write("#!/bin/sh\necho dummy\n")
        os.chmod(p, 0o700)
        return p
    monkeypatch.setattr(rtb, "_write_askpass", fake_askpass)
    try:
        rtb.revert()
    except Exception:
        pass  # expected because fake_run does nothing useful for push
    # Token must never appear in any command argv
    for cmd in captured:
        joined = " ".join(str(c) for c in cmd)
        assert sentinel not in joined, f"token leaked in argv: {joined}"


if __name__ == "__main__":
    pytest.main([__file__, "-q"])


def test_lora_cache_path_includes_revision(tmp_path, monkeypatch):
    from app import _lora_cache_path
    # _lora_cache_path reads _LORA_CACHE_DIR as a module attribute each call.
    # Patch the source module attribute, not the re-export on `app`.
    monkeypatch.setattr("lora.cache._LORA_CACHE_DIR", str(tmp_path))
    p1 = _lora_cache_path("org/repo", "a.safetensors", revision="abc123")
    p2 = _lora_cache_path("org/repo", "a.safetensors", revision="def456")
    assert p1 != p2
    assert "rev_abc123" in p1
    assert "rev_def456" in p2


def test_cache_quota_rejects_when_full(tmp_path, monkeypatch):
    import app
    monkeypatch.setattr("lora.cache._LORA_CACHE_DIR", str(tmp_path))
    monkeypatch.setattr("lora.cache._H3_MAX_LORA_CACHE_BYTES", 1000)
    # Simulate a live in-memory LoRA so eviction cannot free this file.
    live_key = ("live_org/live_repo", "live.safetensors")
    app._CUSTOM_LoRA[live_key] = [("fake", None, None)]
    live_name = "live_org__live_repo__live.safetensors"
    filler = tmp_path / live_name
    filler.write_bytes(b"x" * 900)
    try:
        with pytest.raises(Exception) as ei:
            app._evict_lora_cache_for_quota(200)
        msg = str(ei.value).lower()
        assert "quota" in msg or "exceeded" in msg
    finally:
        app._CUSTOM_LoRA.pop(live_key, None)


def test_registry_no_auto_on_download_logic():
    """Sanity: _download_custom_lora source no longer calls _save after load path for auto-publish.
    Structural check via source inspection (CPU-only)."""
    from pathlib import Path
    src = Path(__file__).resolve().parents[1].joinpath("lora").joinpath("registry.py").read_text()
    # The download function body must mention "not auto-published"
    assert "not auto-published" in src
    assert "Use *Add to community*" in src or "explicit" in src.lower()


def test_registry_validate_rejects_long_description():
    import app
    with pytest.raises(Exception):
        app._validate_registry_entry({
            "repo": "a/b",
            "file": "x.safetensors",
            "description": "z" * 600,
        })


def test_registry_validate_ok():
    import app
    e = app._validate_registry_entry({
        "repo": "owner/name",
        "file": "lora.safetensors",
        "description": "ok",
        "tags": ["a", "b"],
        "scale": 1.5,
        "revision": "deadbeef",
    })
    assert e["repo"] == "owner/name"
    assert e["revision"] == "deadbeef"
    assert e["scale"] == 1.5


def test_registry_atomic_save(tmp_path, monkeypatch):
    import app, json
    reg = tmp_path / "community_loras.json"
    monkeypatch.setattr("lora.registry._COMMUNITY_LORA_PATH", str(reg))
    entries = [
        {"repo": "o/r", "file": "a.safetensors", "url": "https://huggingface.co/o/r/resolve/main/a.safetensors",
         "scale": 1.0, "trigger": "", "description": "", "tags": []},
    ]
    app._save_community_registry(entries)
    assert reg.exists()
    data = json.loads(reg.read_text())
    assert len(data) == 1
    assert data[0]["repo"] == "o/r"


def test_max_lora_bytes_constant():
    import app
    assert app._H3_MAX_LORA_BYTES > 0
    assert app._H3_MAX_LORA_CACHE_BYTES >= app._H3_MAX_LORA_BYTES


# ── Review follow-ups: concurrency invariants, path roots, ContextVar ─────────

def test_task_ctx_isolated_across_threads():
    """ContextVar must not leak task A into task B across concurrent handlers."""
    import app
    import threading

    results = {}
    barrier = threading.Barrier(2)

    def worker(name, task):
        token = app._TASK_CTX.set(task)
        try:
            barrier.wait(timeout=5)
            # After both have set, each must still see its own task
            seen = app._TASK_CTX.get()
            results[name] = seen.task_id if seen else None
        finally:
            app._TASK_CTX.reset(token)

    a = app.GenerationTask.create()
    b = app.GenerationTask.create()
    t1 = threading.Thread(target=worker, args=("a", a))
    t2 = threading.Thread(target=worker, args=("b", b))
    t1.start(); t2.start()
    t1.join(timeout=5); t2.join(timeout=5)
    assert results.get("a") == a.task_id
    assert results.get("b") == b.task_id
    assert a.task_id != b.task_id


def test_request_context_release_clears_applied_tracking():
    """_RequestContext.release must clear tracked adapters even without a real model."""
    import app
    import request as req_mod
    # Patch remove helpers in the module that holds the references used by
    # _RequestContext.release() (request.py). Patching `lora.apply._remove_*`
    # or `app._remove_*` doesn't take effect because request.py imported the
    # originals by name at module load time.
    calls = []
    monkey_remove_custom = lambda key, scale: calls.append(("custom", key, scale))
    monkey_remove_nsfw = lambda key, scale=None: calls.append(("nsfw", key, scale))
    monkey_remove_turbo = lambda: calls.append(("turbo",))
    orig_c = req_mod._remove_custom
    orig_n = req_mod._remove_nsfw
    orig_t = req_mod._remove_turbo
    req_mod._remove_custom = monkey_remove_custom
    req_mod._remove_nsfw = monkey_remove_nsfw
    req_mod._remove_turbo = monkey_remove_turbo
    try:
        task = app.GenerationTask.create()
        ctx = app._RequestContext(task=task)
        ctx.acquire_turbo()
        ctx.acquire_nsfw_lora("nsfw_innie", 1.0)
        ctx.acquire_custom_lora(("org/repo", "f.safetensors"), 0.8)
        assert ctx._turbo is True
        assert len(ctx._nsfw) == 1
        assert len(ctx._custom) == 1
        ctx.release()
        assert ctx._turbo is False
        assert ctx._nsfw == []
        assert ctx._custom == []
        # reverse order: custom, nsfw, turbo
        kinds = [c[0] for c in calls]
        assert kinds == ["custom", "nsfw", "turbo"]
    finally:
        req_mod._remove_custom = orig_c
        req_mod._remove_nsfw = orig_n
        req_mod._remove_turbo = orig_t


def test_mutation_lock_serializes_critical_section():
    """Two threads cannot interleave the locked critical section."""
    import app
    import threading
    order = []
    barrier = threading.Barrier(2)

    def critical(name):
        barrier.wait(timeout=5)
        with app._MODEL_MUTATION_LOCK:
            order.append(f"{name}-enter")
            # Hold the lock briefly so the other thread would interleave if unlocked
            time.sleep(0.05)
            order.append(f"{name}-exit")

    t1 = threading.Thread(target=critical, args=("A",))
    t2 = threading.Thread(target=critical, args=("B",))
    t1.start(); t2.start()
    t1.join(timeout=5); t2.join(timeout=5)
    assert len(order) == 4
    # Nested pairs: enter/exit of one must complete before the other enters
    assert order[0].endswith("-enter") and order[1].endswith("-exit")
    assert order[0][:1] == order[1][:1]
    assert order[2].endswith("-enter") and order[3].endswith("-exit")
    assert order[2][:1] == order[3][:1]


def test_export_rejects_bare_tmpdir(tmp_path, monkeypatch):
    """Bare system tempdir is no longer an allowed export root."""
    import app
    monkeypatch.setattr(app, "_TASK_TEMP_ROOT", str(tmp_path / "tasks"))
    (tmp_path / "tasks").mkdir()
    # A file directly under the system tempdir (outside managed roots) must fail
    outside = Path(tempfile.gettempdir()) / f"h3-export-probe-{uuid.uuid4().hex}.mp4"
    try:
        outside.write_bytes(b"not-a-video")
        with pytest.raises(Exception):
            app.export_last_frame(str(outside))
    finally:
        try:
            outside.unlink()
        except OSError:
            pass


def test_export_rejects_symlink_escape(tmp_path, monkeypatch):
    """Symlink inside an allowed root that points outside must be rejected after realpath."""
    import app
    tasks = tmp_path / "tasks"
    tasks.mkdir()
    monkeypatch.setattr(app, "_TASK_TEMP_ROOT", str(tasks))
    # Create a real file outside, symlink inside the allowed root
    outside = tmp_path / "secret.txt"
    outside.write_text("secret")
    link = tasks / "video.mp4"
    link.symlink_to(outside)
    # realpath(link) == outside, which is not under tasks or h3-outputs
    with pytest.raises(Exception):
        app.export_last_frame(str(link))


def test_write_askpass_has_no_token_and_is_private(tmp_path, monkeypatch):
    import revert_to_baseline as rtb
    sentinel = "hf_SENTINEL_TOKEN_DO_NOT_LEAK_xyz"
    monkeypatch.setenv("HF_TOKEN", sentinel)
    path = rtb._write_askpass(str(tmp_path))
    content = Path(path).read_text()
    assert sentinel not in content
    assert "HF_TOKEN" in content  # reads from env, does not embed value
    mode = Path(path).stat().st_mode & 0o777
    assert mode == 0o700


def test_registry_concurrent_saves_preserve_valid_json(tmp_path, monkeypatch):
    import app, json, threading
    reg = tmp_path / "community_loras.json"
    monkeypatch.setattr("lora.registry._COMMUNITY_LORA_PATH", str(reg))

    def writer(i):
        entries = [{
            "repo": f"owner/r{i}",
            "file": f"l{i}.safetensors",
            "url": f"https://huggingface.co/owner/r{i}/resolve/main/l{i}.safetensors",
            "scale": 1.0, "trigger": "", "description": f"d{i}", "tags": [],
        }]
        # Merge with whatever is already on disk
        existing = app._load_community_registry()
        existing.extend(entries)
        app._save_community_registry(existing)

    threads = [threading.Thread(target=writer, args=(i,)) for i in range(8)]
    for t in threads:
        t.start()
    for t in threads:
        t.join(timeout=5)
    data = json.loads(reg.read_text())
    assert isinstance(data, list)
    # All entries must be valid structure
    for e in data:
        assert "repo" in e and "file" in e
    # At least one entry survived; with last-writer-wins we may not have all 8
    assert len(data) >= 1


def test_parse_civitai_model_page():
    import app
    mid, vid, host, is_arch = app._parse_civitai_link("https://civitai.com/models/12345")
    assert mid == 12345 and vid is None and host == "civitai.com" and is_arch is False
    mid, vid, host, is_arch = app._parse_civitai_link("https://civitai.com/models/12345?modelVersionId=67890")
    assert mid == 12345 and vid == 67890 and not is_arch
    mid, vid, host, is_arch = app._parse_civitai_link("https://civitai.com/api/download/models/999")
    assert mid is None and vid == 999
    # Official .red front door
    mid, vid, host, is_arch = app._parse_civitai_link("https://civitai.red/models/12345?modelVersionId=67890")
    assert mid == 12345 and vid == 67890 and host == "civitai.red" and not is_arch
    mid, vid, host, is_arch = app._parse_civitai_link("https://www.civitai.red/api/download/models/42")
    assert mid is None and vid == 42 and "civitai.red" in host


def test_parse_civitai_rejects_non_civitai():
    import app
    with pytest.raises(Exception):
        app._parse_civitai_link("https://huggingface.co/owner/repo")


def test_parse_civarchive_accepted():
    import app
    mid, vid, host, is_arch = app._parse_civitai_link(
        "https://civitaiarchive.com/models/179979?modelVersionId=201993"
    )
    assert mid == 179979 and vid == 201993 and is_arch is True
    assert "civitaiarchive.com" in host
    mid, vid, host, is_arch = app._parse_civitai_link("https://civarchive.com/models/42")
    assert mid == 42 and vid is None and is_arch is True


def test_civitai_cache_path_stable(tmp_path, monkeypatch):
    import app
    monkeypatch.setattr("lora.cache._LORA_CACHE_DIR", str(tmp_path))
    p1 = app._civitai_cache_path(42, "foo.safetensors")
    p2 = app._civitai_cache_path(42, "foo.safetensors")
    p3 = app._civitai_cache_path(43, "foo.safetensors")
    assert p1 == p2
    assert p1 != p3
    assert "civitai__v42" in p1


def test_civitai_headers_no_token_leak(monkeypatch):
    import app
    monkeypatch.delenv("CIVITAI_API_TOKEN", raising=False)
    monkeypatch.delenv("CIVITAI_API_KEY", raising=False)
    h = app._civitai_headers()
    assert "Authorization" not in h
    monkeypatch.setenv("CIVITAI_API_TOKEN", "secret-token-xyz")
    h2 = app._civitai_headers()
    assert h2.get("Authorization") == "Bearer secret-token-xyz"


def test_parse_civitai_rejects_host_confusion():
    import app
    for bad in (
        "https://civitai.com.evil.com/models/1",
        "https://evilcivitai.com/models/1",
        "https://example.com/models/1",
    ):
        with pytest.raises(Exception):
            app._parse_civitai_link(bad)
    # Scheme is enforced at download time, not parse time.


def test_assert_download_url_allowed_https_and_host():
    import app
    ok = app._assert_download_url_allowed(
        "https://civitai.com/api/download/models/1",
        allowed_hosts=app._CIVITAI_DOWNLOAD_HOSTS,
    )
    assert ok.startswith("https://civitai.com/")
    with pytest.raises(Exception):
        app._assert_download_url_allowed(
            "http://civitai.com/api/download/models/1",
            allowed_hosts=app._CIVITAI_DOWNLOAD_HOSTS,
        )
    with pytest.raises(Exception):
        app._assert_download_url_allowed(
            "https://evil.example/payload.bin",
            allowed_hosts=app._CIVITAI_DOWNLOAD_HOSTS,
        )


def test_resolve_civitai_rewrites_download_url(monkeypatch):
    """API-supplied downloadUrl host must not be trusted."""
    import app
    import types

    class FakeResp:
        status_code = 200
        def json(self):
            return {
                "id": 99,
                "modelId": 10,
                "name": "v1",
                "model": {"name": "TestLoRA"},
                "files": [
                    {
                        "name": "test.safetensors",
                        "primary": True,
                        "sizeKB": 100,
                        "downloadUrl": "https://evil.example/steal",
                        "hashes": {"SHA256": "abc"},
                    }
                ],
                "trainedWords": ["tok"],
            }

    def fake_get(url, headers=None, timeout=None):
        assert "evil" not in url
        return FakeResp()

    monkeypatch.setattr(app, "_civitai_headers", lambda: {})
    import requests
    monkeypatch.setattr(requests, "get", fake_get)
    # Also need requests imported inside function - patch requests module used by app
    monkeypatch.setitem(__import__("sys").modules, "requests", types.SimpleNamespace(get=fake_get, RequestException=Exception))

    meta = app._resolve_civitai_version(None, 99, host="civitai.com")
    assert "evil" not in meta["download_url"]
    assert meta["download_url"].startswith("https://civitai.com/api/download/models/")
    assert meta["filename"] == "test.safetensors"