tests/test_runtime.py

"""可复现的运行时边界测试，不联网，不需要 Blender 或 GPU。

安装测试用本地压缩包替代下载；子进程测试真实启动临时 Python 可执行
程序验证环境和进程树管理，不代表真实 Cycles / ZeroGPU 已通过测试。
"""

from __future__ import annotations

import hashlib
import io
import json
import os
import sys
import tarfile
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import pytest

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import blender_runtime as runtime


pytestmark = pytest.mark.skipif(sys.platform != "linux", reason="目标运行时是 Hugging Face Linux Space")
ARCHIVE_ROOT = "blender-5.2.2-linux-x64"
ARCHIVE_NAME = ARCHIVE_ROOT + ".tar.xz"


def make_archive(path: Path, entries: list[tuple[str, bytes | None, str | None]]) -> Path:
    """None 内容表示目录；第三列非空表示符号链接。"""
    with tarfile.open(path, "w:xz") as archive:
        for name, data, link in entries:
            member = tarfile.TarInfo(name)
            if link is not None:
                member.type = tarfile.SYMTYPE
                member.linkname = link
                archive.addfile(member)
            elif data is None:
                member.type = tarfile.DIRTYPE
                member.mode = 0o755
                archive.addfile(member)
            else:
                member.mode = 0o755
                member.size = len(data)
                archive.addfile(member, io.BytesIO(data))
    return path


@pytest.fixture
def installation(monkeypatch, tmp_path):
    archive = make_archive(tmp_path / ARCHIVE_NAME, [
        (ARCHIVE_ROOT, None, None),
        (ARCHIVE_ROOT + "/blender", b"#!/bin/sh\nexit 0\n", None),
    ])
    checksum = tmp_path / "checksums.sha256"
    digest = hashlib.sha256(archive.read_bytes()).hexdigest()
    checksum.write_text(f"{digest}  {ARCHIVE_NAME}\n", encoding="utf-8")
    cache = tmp_path / "cache"
    monkeypatch.delenv("BLENDER_BIN", raising=False)
    monkeypatch.setenv("BLENDER_VERSION", "5.2.2")
    monkeypatch.setenv("BLENDER_CACHE", str(cache))
    # 测试使用独立目录，不能误用部署仓库随附的真实大安装包。
    monkeypatch.setattr(runtime, "__file__", str(tmp_path / "blender_runtime.py"))
    monkeypatch.setattr(runtime.shutil, "which", lambda _: None)
    calls: list[str] = []

    def local_download(url: str, destination: Path, maximum: int) -> None:
        calls.append(url)
        content = checksum.read_bytes() if url.endswith(".sha256") else archive.read_bytes()
        assert len(content) < maximum
        destination.write_bytes(content)

    monkeypatch.setattr(runtime, "_download", local_download)
    return archive, checksum, cache, calls, digest


def test_download_checksum_is_exact_and_mismatch_stops_installation(installation):
    archive, checksum, cache, calls, digest = installation
    assert runtime._expected_digest(checksum, ARCHIVE_NAME) == digest
    with pytest.raises(RuntimeError, match="唯一匹配"):
        runtime._expected_digest(checksum, "wrong-name.tar.xz")
    # 模拟下载包损坏：官方清单保持不变，下载内容被更改。
    archive.write_bytes(archive.read_bytes() + b"corrupted-download")
    with pytest.raises(RuntimeError, match="SHA-256 校验失败"):
        runtime.ensure_blender()
    assert len(calls) == 2
    assert not (cache / "blender-5.2.2" / "blender").exists()


def test_tar_rejects_traversal_and_absolute_paths(tmp_path):
    for index, name in enumerate([
        ARCHIVE_ROOT + "/../../outside", "/tmp/outside", ARCHIVE_ROOT + "/../outside",
    ]):
        archive = make_archive(tmp_path / f"bad-{index}.tar.xz", [(name, b"unsafe", None)])
        with pytest.raises(RuntimeError, match="不安全路径"):
            runtime._safe_extract(archive, tmp_path / "output", ARCHIVE_ROOT)
    assert not (tmp_path / "outside").exists()


def test_tar_accepts_internal_symlinks_and_rejects_escaping_links(tmp_path):
    archive = make_archive(tmp_path / "good.tar.xz", [
        (ARCHIVE_ROOT, None, None),
        (ARCHIVE_ROOT + "/blender", b"executable", None),
        (ARCHIVE_ROOT + "/linked", None, "blender"),
    ])
    extracted = runtime._safe_extract(archive, tmp_path / "good", ARCHIVE_ROOT)
    assert (extracted / "linked").is_symlink()
    assert (extracted / "linked").read_bytes() == b"executable"
    bad = make_archive(tmp_path / "escape.tar.xz", [
        (ARCHIVE_ROOT + "/linked", None, "../../outside"),
    ])
    with pytest.raises(RuntimeError, match="目录之外"):
        runtime._safe_extract(bad, tmp_path / "bad", ARCHIVE_ROOT)


def test_concurrent_installation_downloads_each_file_once(installation):
    _, _, cache, calls, _ = installation
    # Future.result() 会把线程中的异常带回测试主线程，避免漏报。
    with ThreadPoolExecutor(max_workers=5) as executor:
        futures = [executor.submit(runtime.ensure_blender) for _ in range(5)]
        binaries = [future.result(timeout=10) for future in futures]
    assert len(set(binaries)) == 1
    assert binaries[0] == cache / "blender-5.2.2" / "blender"
    assert os.access(binaries[0], os.X_OK)
    assert len(calls) == 2
    assert sum(url.endswith(".sha256") for url in calls) == 1
    assert sum(url.endswith(".tar.xz") for url in calls) == 1
    assert runtime.ensure_blender() == binaries[0]
    assert len(calls) == 2


def test_bundled_archive_installs_without_network(installation, monkeypatch, tmp_path):
    archive, checksum, cache, calls, _ = installation
    vendor = tmp_path / "vendor"
    vendor.mkdir()
    (vendor / ARCHIVE_NAME).write_bytes(archive.read_bytes())
    (vendor / "blender-5.2.2.sha256").write_bytes(checksum.read_bytes())
    monkeypatch.setattr(runtime, "__file__", str(tmp_path / "blender_runtime.py"))
    binary = runtime.ensure_blender()
    assert binary == cache / "blender-5.2.2" / "blender"
    assert os.access(binary, os.X_OK)
    assert calls == []


def test_corrupt_bundled_archive_is_rejected_without_network(installation, monkeypatch, tmp_path):
    archive, checksum, cache, calls, _ = installation
    vendor = tmp_path / "vendor"
    vendor.mkdir()
    (vendor / ARCHIVE_NAME).write_bytes(archive.read_bytes() + b"damaged")
    (vendor / "blender-5.2.2.sha256").write_bytes(checksum.read_bytes())
    monkeypatch.setattr(runtime, "__file__", str(tmp_path / "blender_runtime.py"))
    with pytest.raises(RuntimeError, match="SHA-256 校验失败"):
        runtime.ensure_blender()
    assert calls == []
    assert not (cache / "blender-5.2.2").exists()


@pytest.fixture
def subprocess_runtime(monkeypatch, tmp_path):
    fake = tmp_path / "fake-blender"
    fake.write_text("#!" + sys.executable + "\n" + r'''
import json, os, subprocess, sys, time
from pathlib import Path
task = json.loads(Path(sys.argv[-1]).read_text(encoding="utf-8"))
if task.get("delay"):
    child = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(60)"])
    Path(task["child_pid"]).write_text(str(child.pid), encoding="utf-8")
    time.sleep(60)
else:
    Path(task["result_path"]).write_text(json.dumps({
        "ok": True,
        "secret_present": "HF_TOKEN" in os.environ or "AWS_SECRET_ACCESS_KEY" in os.environ,
        "cuda": os.environ.get("CUDA_VISIBLE_DEVICES"),
        "threads": os.environ.get("OMP_NUM_THREADS"),
        "autoexec_disabled": "--disable-autoexec" in sys.argv,
    }), encoding="utf-8")
''', encoding="utf-8")
    fake.chmod(0o755)
    monkeypatch.setattr(runtime, "_WORKER", fake)
    monkeypatch.setattr(runtime, "ensure_blender", lambda: fake)
    monkeypatch.setenv("BLENDER_THREADS", "4")
    return tmp_path


def test_real_subprocess_scrubs_secrets_and_preserves_cuda(subprocess_runtime, monkeypatch):
    root = subprocess_runtime
    monkeypatch.setenv("HF_TOKEN", "test-only-never-log")
    monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "test-only-never-log")
    monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "test-gpu-uuid")
    result = runtime.run_task({"result_path": str(root / "result.json")}, 5, root / "render.log")
    assert result == {
        "ok": True, "secret_present": False, "cuda": "test-gpu-uuid",
        "threads": "4", "autoexec_disabled": True,
    }
    assert not list(root.glob(".render-task-*.json"))


def assert_descendant_stopped(pid_file: Path) -> None:
    assert pid_file.exists(), "测试替身未成功创建子进程，无法验证进程树终止。"
    pid = int(pid_file.read_text(encoding="utf-8"))
    stat = Path(f"/proc/{pid}/stat")
    # 被结束的孤儿子进程可能暂时等待容器 init 回收；Z 状态已不再执行。
    deadline = time.monotonic() + 2
    while stat.exists() and time.monotonic() < deadline:
        try:
            if stat.read_text().rsplit(")", 1)[1].split()[0] == "Z":
                return
        except FileNotFoundError:
            return
        time.sleep(0.02)
    assert not stat.exists(), f"后代进程 {pid} 仍在运行。"


def test_timeout_terminates_real_descendant_process(subprocess_runtime):
    root = subprocess_runtime
    pid_file = root / "timeout-child.pid"
    task = {"result_path": str(root / "timeout.json"), "delay": True, "child_pid": str(pid_file)}
    with pytest.raises(TimeoutError, match="超过"):
        runtime.run_task(task, 1, root / "timeout.log")
    assert_descendant_stopped(pid_file)
    assert not list(root.glob(".render-task-*.json"))


def test_cancellation_terminates_real_descendant_process(subprocess_runtime):
    root = subprocess_runtime
    pid_file = root / "cancel-child.pid"
    cancel = root / "cancel.flag"
    task = {"result_path": str(root / "cancel.json"), "delay": True, "child_pid": str(pid_file)}

    def cancel_when_child_started():
        deadline = time.monotonic() + 5
        while not pid_file.exists() and time.monotonic() < deadline:
            time.sleep(0.02)
        cancel.touch()

    thread = threading.Thread(target=cancel_when_child_started, daemon=True)
    thread.start()
    try:
        with pytest.raises(RuntimeError, match="取消"):
            runtime.run_task(task, 8, root / "cancel.log", cancel)
    finally:
        thread.join(timeout=6)
    assert_descendant_stopped(pid_file)
    assert not list(root.glob(".render-task-*.json"))