"""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"