mmh3testing / tests /test_task_isolation.py
jhickman51's picture
Merge upstream MiniMax H3 updates and preserve beta5 checkpoint
b0eb306
Raw History Blame Contribute Delete
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"