Spaces:
Running on Zero
Running on Zero
Download tests/test_task_isolation.py from jickman125/mmh3testing: direct link, hf CLI and curl.
- Browser
- Download file 18.8 kB
-
https://huggingface.co/spaces/jickman125/mmh3testing/resolve/main/tests/test_task_isolation.py
- Command line
-
hf download hf://spaces/jickman125/mmh3testing/tests/test_task_isolation.py
-
curl -L -o test_task_isolation.py https://huggingface.co/spaces/jickman125/mmh3testing/resolve/main/tests/test_task_isolation.py
18.8 kB
| """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" | |