wov-api/tests/test_node_manager.py

"""NodeManager 单元测试。

覆盖命令解析、环境变量注入、CUDA 路径处理、进程启动失败分支、实例复用、
调用错误处理、回收线程以及停止清理等真实生命周期路径。
"""

import os
import subprocess
import sys
import time
from datetime import datetime, timedelta, timezone
from pathlib import Path

import pytest

from app.db import Database
from app.node_manager import (
    NodeManager,
    NodeRuntime,
    _cuda_library_dirs,
    _idle_seconds,
    _now_iso,
    _with_cuda_library_path,
)
from wov_sdk.models import InvokeRequest, NodeManifest

WORKSPACE = Path(__file__).resolve().parent.parent.parent
# 从真实节点仓库加载 Echo manifest，测试使用真实协议数据。
_BASE_ECHO_MANIFEST = NodeManifest.load(WORKSPACE / "wov-node-echo" / "node.manifest.json")
FIXTURES = Path(__file__).resolve().parent / "fixtures"


@pytest.fixture
def db(tmp_path) -> Database:
    return Database(tmp_path / "wov.db")


def _register(db: Database, manifest: NodeManifest) -> None:
    """把 manifest 写入测试数据库，模拟注册节点。"""
    db.upsert_node(manifest)


def echo_manifest() -> NodeManifest:
    """返回 Echo 节点 manifest 的副本，避免测试间共享可变对象。"""
    return NodeManifest.from_dict(_BASE_ECHO_MANIFEST.to_dict())


def test_time_helpers() -> None:
    """验证时间工具对正常、空与非法输入的返回。"""
    now = _now_iso()
    assert datetime.fromisoformat(now)
    assert _idle_seconds(now) < 0.1
    old = (datetime.now(timezone.utc) - timedelta(seconds=10)).isoformat()
    assert _idle_seconds(old) >= 9
    assert _idle_seconds("") == 0
    assert _idle_seconds("invalid") == 0


def test_resolve_command_and_env(db: Database) -> None:
    """验证 python 命令解析回退与固定命令保持原样。"""
    manager = NodeManager(db)
    python_command = NodeManifest(
        id="x",
        name="x",
        version="1",
        capability="x",
        command=["python", "-m", "x"],
    )
    assert manager._resolve_command(python_command)[0] == sys.executable

    fixed_command = NodeManifest(
        id="x",
        name="x",
        version="1",
        capability="x",
        command=["node", "index.js"],
    )
    assert manager._resolve_command(fixed_command) == ["node", "index.js"]

    env = manager._node_env(python_command)
    assert "wov-sdk" in env["PYTHONPATH"]


def test_resolve_command_prefers_node_venv(db: Database, tmp_path, monkeypatch) -> None:
    """验证 Windows 与 Linux 虚拟环境解释器优先级。"""
    monkeypatch.setattr("app.node_manager.WORKSPACE_ROOT", tmp_path)
    repo = tmp_path / "wov-node-demo"
    python_exe = repo / ".venv" / "Scripts" / "python.exe"
    python_exe.parent.mkdir(parents=True)
    python_exe.write_bytes(b"")
    manifest = NodeManifest(
        id="x",
        name="x",
        version="1",
        capability="x",
        command=["python", "-m", "x"],
        repo_dir="wov-node-demo",
    )
    manager = NodeManager(db)
    assert manager._resolve_command(manifest)[0] == str(python_exe)

    python_exe.unlink()
    unix_python = repo / ".venv" / "bin" / "python"
    unix_python.parent.mkdir(parents=True)
    unix_python.write_bytes(b"")
    assert manager._resolve_command(manifest)[0] == str(unix_python)


def test_cuda_library_dirs_finds_unix_nvidia_libs(tmp_path) -> None:
    """验证 Linux site-packages 下的 cublas/cudnn lib 目录被发现。"""
    cublas = (
        tmp_path / ".venv" / "lib" / "python3.12" / "site-packages" / "nvidia" / "cublas" / "lib"
    )
    cudnn = (
        tmp_path / ".venv" / "lib" / "python3.12" / "site-packages" / "nvidia" / "cudnn" / "lib"
    )
    cublas.mkdir(parents=True)
    cudnn.mkdir(parents=True)

    dirs = _cuda_library_dirs(tmp_path)

    assert cublas in dirs
    assert cudnn in dirs


def test_cuda_library_dirs_finds_windows_nvidia_bins(tmp_path) -> None:
    """验证 Windows Lib/site-packages 下的 bin 目录被发现。"""
    cublas = tmp_path / ".venv" / "Lib" / "site-packages" / "nvidia" / "cublas" / "bin"
    cublas.mkdir(parents=True)

    dirs = _cuda_library_dirs(tmp_path)

    assert cublas in dirs


def test_with_cuda_library_path_prepends_on_posix(monkeypatch) -> None:
    """验证 Linux 下 CUDA 目录插入 LD_LIBRARY_PATH 开头。"""
    monkeypatch.setattr(os, "name", "posix")
    monkeypatch.setattr(
        "app.node_manager._cuda_library_dirs",
        lambda repo_dir: ["/opt/nvidia/cublas/lib", "/opt/nvidia/cudnn/lib"],
    )

    env = _with_cuda_library_path({"LD_LIBRARY_PATH": "/usr/lib/foo"}, Path("/repo"))

    entries = env["LD_LIBRARY_PATH"].split(os.pathsep)
    assert entries[:2] == ["/opt/nvidia/cublas/lib", "/opt/nvidia/cudnn/lib"]
    assert entries[-1] == "/usr/lib/foo"


def test_with_cuda_library_path_uses_path_on_windows(monkeypatch) -> None:
    """验证 Windows 下 CUDA 目录插入 PATH 开头。"""
    monkeypatch.setattr(os, "name", "nt")
    monkeypatch.setattr(
        "app.node_manager._cuda_library_dirs",
        lambda repo_dir: ["/opt/nvidia/cublas/bin"],
    )

    env = _with_cuda_library_path({"PATH": "/usr/local/bin"}, Path("/repo"))

    entries = env["PATH"].split(os.pathsep)
    assert entries[0] == "/opt/nvidia/cublas/bin"
    assert entries[-1] == "/usr/local/bin"


def test_with_cuda_library_path_without_existing(monkeypatch) -> None:
    """验证环境变量原本不存在时直接创建。"""
    monkeypatch.setattr(os, "name", "posix")
    monkeypatch.setattr(
        "app.node_manager._cuda_library_dirs",
        lambda repo_dir: ["/opt/nvidia/cublas/lib"],
    )

    env = _with_cuda_library_path({}, Path("/repo"))

    assert env["LD_LIBRARY_PATH"] == "/opt/nvidia/cublas/lib"


def test_with_cuda_library_path_deduplicates(monkeypatch) -> None:
    """验证已存在的目录不会被重复追加。"""
    monkeypatch.setattr(os, "name", "posix")
    monkeypatch.setattr(
        "app.node_manager._cuda_library_dirs",
        lambda repo_dir: ["/opt/nvidia/cublas/lib", "/opt/nvidia/cudnn/lib"],
    )

    env = _with_cuda_library_path(
        {"LD_LIBRARY_PATH": "/opt/nvidia/cublas/lib:/usr/lib/foo"},
        Path("/repo"),
    )

    entries = env["LD_LIBRARY_PATH"].split(os.pathsep)
    assert entries == ["/opt/nvidia/cudnn/lib", "/opt/nvidia/cublas/lib", "/usr/lib/foo"]


def test_with_cuda_library_path_all_present(monkeypatch) -> None:
    """验证所有目录都已存在时环境变量保持不变。"""
    monkeypatch.setattr(os, "name", "posix")
    monkeypatch.setattr(
        "app.node_manager._cuda_library_dirs",
        lambda repo_dir: ["/opt/nvidia/cublas/lib"],
    )
    original = {"LD_LIBRARY_PATH": "/opt/nvidia/cublas/lib:/usr/lib/foo"}

    env = _with_cuda_library_path(dict(original), Path("/repo"))

    assert env["LD_LIBRARY_PATH"] == original["LD_LIBRARY_PATH"]


def test_with_cuda_library_path_without_libs(monkeypatch) -> None:
    """验证没有任何 CUDA 库时环境变量原样返回。"""
    monkeypatch.setattr("app.node_manager._cuda_library_dirs", lambda repo_dir: [])

    env = _with_cuda_library_path({"LD_LIBRARY_PATH": "/usr/lib/foo"}, Path("/repo"))

    assert env == {"LD_LIBRARY_PATH": "/usr/lib/foo"}


def test_node_env_injects_cuda_library_path(db: Database, tmp_path, monkeypatch) -> None:
    """验证 _node_env 会注入 SDK、manifest 变量与 CUDA 路径。"""
    monkeypatch.setattr("app.node_manager.WORKSPACE_ROOT", tmp_path)
    repo = tmp_path / "wov-node-whisper"
    cublas = repo / ".venv" / "lib" / "python3.12" / "site-packages" / "nvidia" / "cublas" / "lib"
    cublas.mkdir(parents=True)
    manifest = NodeManifest(
        id="x",
        name="x",
        version="1",
        capability="x",
        command=["python", "-m", "x"],
        repo_dir="wov-node-whisper",
        env={"WOV_NODE_PORT": "0"},
    )

    env = NodeManager(db)._node_env(manifest)

    assert str(cublas) in env["LD_LIBRARY_PATH"].split(os.pathsep)
    assert env["WOV_NODE_PORT"] == "0"
    assert "wov-sdk" in env["PYTHONPATH"]


def test_acquire_unregistered(db: Database) -> None:
    """验证未注册节点无法获取实例。"""
    manager = NodeManager(db)
    with pytest.raises(ValueError):
        manager.acquire("missing")


def test_start_missing_repo(db: Database) -> None:
    """验证 repo_dir 不存在时启动失败并记录 error。"""
    manifest = echo_manifest()
    manifest.repo_dir = "missing-repo"
    _register(db, manifest)
    manager = NodeManager(db)
    with pytest.raises(RuntimeError, match="repo_dir not found"):
        manager.acquire("echo")
    assert db.list_instances()[0]["status"] == "error"


def test_start_command_not_found(db: Database) -> None:
    """验证命令不存在时启动失败并记录 error。"""
    manifest = echo_manifest()
    manifest.command = ["definitely-not-a-real-wov-command"]
    _register(db, manifest)
    manager = NodeManager(db)
    with pytest.raises(RuntimeError):
        manager.acquire("echo")
    assert db.list_instances()[0]["status"] == "error"


def test_start_no_ready_timeout(db: Database, monkeypatch) -> None:
    """验证超时未打印就绪行时启动失败。"""
    monkeypatch.setattr("app.node_manager.NODE_READY_TIMEOUT_SECONDS", 0.2)
    manifest = echo_manifest()
    manifest.command = ["python", "-c", "import time; time.sleep(0.5)"]
    _register(db, manifest)
    manager = NodeManager(db)
    with pytest.raises(RuntimeError, match="did not report readiness"):
        manager.acquire("echo")
    assert db.list_instances()[0]["status"] == "error"


def test_start_bad_ready_line(db: Database) -> None:
    """验证就绪行缺少端口时启动失败。"""
    manifest = echo_manifest()
    manifest.command = ["python", "-u", str(FIXTURES / "bad_ready_node.py")]
    _register(db, manifest)
    manager = NodeManager(db)
    with pytest.raises(RuntimeError, match="did not report readiness"):
        manager.acquire("echo")
    assert db.list_instances()[0]["status"] == "error"


def test_start_health_check_failure(db: Database) -> None:
    """验证 /health 返回错误状态码时启动失败。"""
    manifest = echo_manifest()
    manifest.command = ["python", "-u", str(FIXTURES / "bad_node.py")]
    _register(db, manifest)
    manager = NodeManager(db)
    with pytest.raises(RuntimeError, match="health check failed"):
        manager.acquire("echo")
    assert db.list_instances()[0]["status"] == "error"


def test_start_health_check_bad_status(db: Database) -> None:
    """验证 /health 返回非 200 状态码时启动失败。"""
    manifest = echo_manifest()
    manifest.command = ["python", "-u", str(FIXTURES / "bad_status_node.py")]
    _register(db, manifest)
    manager = NodeManager(db)
    with pytest.raises(RuntimeError, match="health check returned 201"):
        manager.acquire("echo")
    assert db.list_instances()[0]["status"] == "error"


def test_stderr_is_captured(db: Database) -> None:
    """验证节点 stderr 会被缓存到运行时供诊断。"""
    manifest = echo_manifest()
    manifest.command = ["python", "-u", str(FIXTURES / "stderr_node.py")]
    _register(db, manifest)
    manager = NodeManager(db)
    try:
        runtime, _ = manager.acquire("echo")
        time.sleep(0.2)
        runtime = next(iter(manager._runtimes.values()))
        assert "stderr warning" in runtime.stderr_tail
        manager.release(runtime.instance_id)
    finally:
        manager.shutdown()


def test_acquire_reuses_ready_instance(db: Database) -> None:
    """验证释放后的就绪实例会被再次复用。"""
    _register(db, echo_manifest())
    manager = NodeManager(db)
    try:
        first, first_address = manager.acquire("echo")
        manager.release(first.instance_id)
        second, second_address = manager.acquire("echo")
        assert first.instance_id == second.instance_id
        assert first_address == second_address
        manager.release(second.instance_id)
    finally:
        manager.shutdown()


def test_invoke_success(db: Database, tmp_path) -> None:
    """验证真实节点调用返回完成状态与输出。"""
    _register(db, echo_manifest())
    manager = NodeManager(db)
    try:
        response = manager.invoke(
            "echo",
            InvokeRequest(
                run_id="run_ok",
                node_instance_id="",
                inputs={"text": "real path"},
                output_dir=str(tmp_path),
            ),
        )
        assert response.status == "completed"
        assert response.outputs["text"] == "real path"
    finally:
        manager.shutdown()


def test_invoke_http_error(db: Database, tmp_path) -> None:
    """验证节点返回协议 JSON 错误时透传错误信息。"""
    manifest = echo_manifest()
    manifest.command = ["python", "-u", str(FIXTURES / "failing_node.py")]
    _register(db, manifest)
    manager = NodeManager(db)
    try:
        response = manager.invoke(
            "echo",
            InvokeRequest(
                run_id="run_fail",
                node_instance_id="",
                inputs={},
                output_dir=str(tmp_path),
            ),
        )
        assert response.status == "failed"
        assert response.error == "boom"
    finally:
        manager.shutdown()


def test_invoke_http_error_invalid_json(db: Database, tmp_path) -> None:
    """验证节点返回非法 JSON 时降级为通用错误信息。"""
    manifest = echo_manifest()
    manifest.command = ["python", "-u", str(FIXTURES / "failing_node.py")]
    manifest.env["WOV_FAIL_INVALID_JSON"] = "1"
    _register(db, manifest)
    manager = NodeManager(db)
    try:
        response = manager.invoke(
            "echo",
            InvokeRequest(
                run_id="run_bad_json",
                node_instance_id="",
                inputs={},
                output_dir=str(tmp_path),
            ),
        )
        assert response.status == "failed"
        assert "node returned 500" in response.error
    finally:
        manager.shutdown()


def test_invoke_network_error(db: Database, tmp_path) -> None:
    """验证进程退出导致连接失败时返回 failed。"""
    _register(db, echo_manifest())
    manager = NodeManager(db)
    try:
        runtime, _ = manager.acquire("echo")
        manager.release(runtime.instance_id)
        runtime.process.terminate()
        runtime.process.wait(timeout=5)
        response = manager.invoke(
            "echo",
            InvokeRequest(
                run_id="run_net",
                node_instance_id="",
                inputs={},
                output_dir=str(tmp_path),
            ),
        )
        assert response.status == "failed"
    finally:
        manager.shutdown()


def test_invoke_uses_configured_timeout(db: Database, monkeypatch) -> None:
    """验证调用超时取自配置常量。"""
    _register(db, echo_manifest())
    manager = NodeManager(db)
    captured = {}

    class FakeRuntime:
        instance_id = "ni_x"
        status = "ready"
        busy_count = 0

    class FakeResponse:
        def read(self) -> bytes:
            return b'{"status":"completed","outputs":{"text":"ok"}}'

        def __enter__(self):
            return self

        def __exit__(self, *args) -> bool:
            return False

    def fake_acquire(node_id):
        return FakeRuntime(), "http://127.0.0.1:1"

    def fake_urlopen(request, timeout):
        captured["timeout"] = timeout
        return FakeResponse()

    monkeypatch.setattr("app.node_manager.NODE_INVOKE_TIMEOUT_SECONDS", 123)
    monkeypatch.setattr(manager, "acquire", fake_acquire)
    monkeypatch.setattr("app.node_manager.urllib.request.urlopen", fake_urlopen)

    response = manager.invoke(
        "echo",
        InvokeRequest(
            run_id="run_t",
            node_instance_id="",
            inputs={},
            output_dir=".",
        ),
    )

    assert response.status == "completed"
    assert captured["timeout"] == 123


def test_release_unknown(db: Database) -> None:
    """验证释放未知实例是无害操作。"""
    manager = NodeManager(db)
    manager.release("missing")


def test_stop_unknown_instance(db: Database) -> None:
    """验证停止未知实例是无害操作。"""
    manager = NodeManager(db)
    manager.stop_instance("missing")


def test_stop_all_for_node(db: Database) -> None:
    """验证 stop_all_for_node 会停止该节点全部实例。"""
    _register(db, echo_manifest())
    manager = NodeManager(db)
    runtime, _ = manager.acquire("echo")
    manager.stop_all_for_node("echo")
    assert runtime.status == "stopped"
    assert runtime.instance_id not in manager._runtimes


def test_stop_locked_without_process(db: Database) -> None:
    """验证无进程句柄的实例也能被标记停止。"""
    _register(db, echo_manifest())
    manager = NodeManager(db)
    runtime = NodeRuntime(
        instance_id="ni_static",
        node_id="echo",
        manifest=echo_manifest(),
        status="ready",
    )
    manager._runtimes[runtime.instance_id] = runtime
    manager.stop_instance(runtime.instance_id)
    assert runtime.status == "stopped"
    assert runtime.instance_id not in manager._runtimes


def test_stop_process_already_exited() -> None:
    """验证进程已退出时停止流程直接返回。"""
    class FakeProcess:
        def poll(self):
            return 0

    manager = object.__new__(NodeManager)
    manager._stop_process(FakeProcess())


def test_stop_process_timeout() -> None:
    """验证优雅终止超时后强制 kill。"""
    class FakeProcess:
        def poll(self):
            return None

        def terminate(self):
            self.terminated = True

        def wait(self, timeout):
            if not self.killed:
                raise subprocess.TimeoutExpired("fake", timeout)

        def kill(self):
            self.killed = True

    fake = FakeProcess()
    fake.killed = False
    manager = object.__new__(NodeManager)
    manager._stop_process(fake)
    assert fake.terminated
    assert fake.killed


def test_reaper_recycles_idle(db: Database, tmp_path) -> None:
    """验证空闲超过 TTL 的实例会被回收线程停止。"""
    manifest = echo_manifest()
    manifest.idle_ttl_seconds = 0
    _register(db, manifest)
    manager = NodeManager(db)
    manager.start_reaper()
    try:
        response = manager.invoke(
            "echo",
            InvokeRequest(
                run_id="run_reap",
                node_instance_id="",
                inputs={"text": "reap"},
                output_dir=str(tmp_path),
            ),
        )
        assert response.status == "completed"
        time.sleep(0.4)
        assert db.list_instances()[0]["status"] == "stopped"
    finally:
        manager.shutdown()


def test_reaper_keeps_warm(db: Database, tmp_path) -> None:
    """验证 keep_warm 实例即使空闲也不会被回收。"""
    manifest = echo_manifest()
    manifest.idle_ttl_seconds = 0
    manifest.keep_warm = True
    _register(db, manifest)
    manager = NodeManager(db)
    manager.start_reaper()
    try:
        response = manager.invoke(
            "echo",
            InvokeRequest(
                run_id="run_warm",
                node_instance_id="",
                inputs={"text": "warm"},
                output_dir=str(tmp_path),
            ),
        )
        assert response.status == "completed"
        time.sleep(0.4)
        assert db.list_instances()[0]["status"] == "ready"
    finally:
        manager.shutdown()


def test_reaper_keeps_busy(db: Database) -> None:
    """验证正在被占用的实例不会被回收。"""
    manifest = echo_manifest()
    manifest.idle_ttl_seconds = 0
    _register(db, manifest)
    manager = NodeManager(db)
    manager.start_reaper()
    try:
        runtime, _ = manager.acquire("echo")
        time.sleep(0.4)
        assert runtime.status == "ready"
        manager.release(runtime.instance_id)
    finally:
        manager.shutdown()