vrsub/tests/test_scheduler.py

"""调度器单元测试。

覆盖拓扑排序、任务执行成功/失败分支、输入引用解析、MIME 推断以及
后台轮询线程的启动与停止。节点调用改为进程内注册表直接调用。
"""

import time
from pathlib import Path

import pytest

from wov_app import registry
from wov_app.db import Database
from wov_app.scheduler import WorkflowScheduler, topological_sort
from wov_sdk.models import (
    InvokeResponse,
    NodeManifest,
    WorkflowDefinition,
    WorkflowEdge,
    WorkflowNode,
)

# 单体根目录：tests/ 的上一级。
WORKSPACE = Path(__file__).resolve().parent.parent


def _register_echo() -> None:
    """把内置 echo 节点注册到进程内注册表。"""
    from nodes.echo import invoke

    registry.register(NodeManifest.load(str(WORKSPACE / "manifests" / "echo.json")), invoke)


def _db(tmp_path) -> Database:
    """在临时目录创建独立数据库。"""
    return Database(tmp_path / "wov.db")


def _echo_definition() -> WorkflowDefinition:
    """构造引用 Echo 节点的单步骤工作流定义。"""
    return WorkflowDefinition(
        name="echo-flow",
        version=1,
        nodes=[
            WorkflowNode(
                id="step",
                node_type="echo",
                inputs={"file_uri": "input.video_uri"},
            )
        ],
        edges=[],
        entry_inputs={"video_uri": "file"},
        final_outputs={"result": "step.file_uri"},
    )


def test_topological_sort() -> None:
    """验证 DAG 排序保持依赖顺序，并拒绝环与未知边。"""
    definition = WorkflowDefinition(
        name="dag",
        version=1,
        nodes=[
            WorkflowNode(id="a", node_type="x"),
            WorkflowNode(id="b", node_type="x"),
            WorkflowNode(id="c", node_type="x"),
        ],
        edges=[
            WorkflowEdge(from_node="a", to_node="b"),
            WorkflowEdge(from_node="a", to_node="c"),
        ],
    )
    order = topological_sort(definition)
    assert order.index("a") < order.index("b")
    assert order.index("a") < order.index("c")

    cycle = WorkflowDefinition(
        name="cycle",
        version=1,
        nodes=[
            WorkflowNode(id="a", node_type="x"),
            WorkflowNode(id="b", node_type="x"),
        ],
        edges=[
            WorkflowEdge(from_node="a", to_node="b"),
            WorkflowEdge(from_node="b", to_node="a"),
        ],
    )
    with pytest.raises(ValueError, match="cycle"):
        topological_sort(cycle)

    with pytest.raises(ValueError, match="unknown edge"):
        topological_sort(
            WorkflowDefinition(
                name="bad",
                version=1,
                nodes=[WorkflowNode(id="a", node_type="x")],
                edges=[WorkflowEdge(from_node="a", to_node="missing")],
            )
        )


def test_execute_echo_workflow(tmp_path) -> None:
    """验证排队任务可被完整执行并登记全部产物。"""
    db = _db(tmp_path)
    input_file = tmp_path / "input.txt"
    input_file.write_text("hello scheduler", encoding="utf-8")
    _register_echo()
    db.upsert_workflow({"id": "flow", "name": "Flow", "published": 1, "latest_version": 1})
    db.create_workflow_version("flow", 1, _echo_definition().to_dict())
    now = "2026-01-01T00:00:00+00:00"
    db.create_run(
        {
            "id": "run_1",
            "workflow_id": "flow",
            "workflow_version": 1,
            "status": "QUEUED",
            "progress": 0,
            "input_uri": str(input_file),
            "created_at": now,
            "updated_at": now,
        }
    )

    scheduler = WorkflowScheduler(db, tmp_path / "storage")
    scheduler.execute_run("run_1")

    run = db.get_run("run_1")
    assert run["status"] == "COMPLETED"
    artifacts = db.list_artifacts("run_1")
    assert {item["name"] for item in artifacts} == {"step.text", "step.file_uri", "result"}


def test_execute_run_missing_workflow(tmp_path, monkeypatch) -> None:
    """验证工作流记录缺失时任务被标记为失败。"""
    db = _db(tmp_path)
    db.upsert_workflow({"id": "flow", "name": "Flow", "published": 1, "latest_version": 1})
    now = "2026-01-01T00:00:00+00:00"
    db.create_run(
        {
            "id": "run_missing",
            "workflow_id": "flow",
            "workflow_version": 1,
            "status": "QUEUED",
            "progress": 0,
            "created_at": now,
            "updated_at": now,
        }
    )
    monkeypatch.setattr(db, "get_workflow", lambda workflow_id: None)
    scheduler = WorkflowScheduler(db, tmp_path / "storage")
    scheduler.execute_run("run_missing")
    assert db.get_run("run_missing")["status"] == "FAILED"


def test_execute_run_missing_version(tmp_path) -> None:
    """验证版本记录缺失时任务被标记为失败。"""
    db = _db(tmp_path)
    db.upsert_workflow({"id": "flow", "name": "Flow", "published": 1, "latest_version": 1})
    now = "2026-01-01T00:00:00+00:00"
    db.create_run(
        {
            "id": "run_version",
            "workflow_id": "flow",
            "workflow_version": 1,
            "status": "QUEUED",
            "progress": 0,
            "created_at": now,
            "updated_at": now,
        }
    )
    scheduler = WorkflowScheduler(db, tmp_path / "storage")
    scheduler.execute_run("run_version")
    assert db.get_run("run_version")["status"] == "FAILED"


def test_execute_run_missing_node(tmp_path) -> None:
    """验证未注册节点被调用时任务失败。"""
    db = _db(tmp_path)
    definition = WorkflowDefinition(
        name="bad",
        version=1,
        nodes=[
            WorkflowNode(
                id="step",
                node_type="missing-node",
                inputs={"text": "input.video_uri"},
            )
        ],
        entry_inputs={"video_uri": "file"},
    )
    db.upsert_workflow({"id": "flow", "name": "Flow", "published": 1, "latest_version": 1})
    db.create_workflow_version("flow", 1, definition.to_dict())
    now = "2026-01-01T00:00:00+00:00"
    db.create_run(
        {
            "id": "run_node",
            "workflow_id": "flow",
            "workflow_version": 1,
            "status": "QUEUED",
            "progress": 0,
            "input_uri": str(tmp_path / "in.txt"),
            "created_at": now,
            "updated_at": now,
        }
    )
    scheduler = WorkflowScheduler(db, tmp_path / "storage")
    scheduler.execute_run("run_node")
    assert db.get_run("run_node")["status"] == "FAILED"


def test_resolve_ref_and_mime(tmp_path) -> None:
    """验证输入引用解析、MIME 推断与文件大小读取。"""
    db = _db(tmp_path)
    scheduler = WorkflowScheduler(db, tmp_path / "storage")
    assert scheduler._resolve_ref("input.video", "in.mp4", {}) == "in.mp4"
    assert (
        scheduler._resolve_ref(
            "a.out", None, {"a": {"out": "result.txt"}}
        )
        == "result.txt"
    )
    assert scheduler._resolve_ref("a.out", None, {}) is None
    assert scheduler._resolve_ref("nodot", "in.mp4", {}) is None
    assert scheduler._mime_type("x.srt") == "application/x-subrip"
    assert scheduler._mime_type("x.ass") == "text/plain"
    assert scheduler._mime_type("x.wav") == "audio/wav"
    assert scheduler._mime_type("x.mp4") == "video/mp4"
    assert scheduler._mime_type("x.txt") == "text/plain"
    assert scheduler._mime_type("x.bin") == "application/octet-stream"
    existing = tmp_path / "existing.txt"
    existing.write_text("x", encoding="utf-8")
    assert scheduler._file_size(str(existing)) == 1
    missing = tmp_path / "missing.bin"
    assert scheduler._file_size(str(missing)) == 0


def test_execute_unknown_or_non_queued_run(tmp_path) -> None:
    """验证未知任务或非排队任务会被忽略。"""
    db = _db(tmp_path)
    db.upsert_workflow({"id": "flow", "name": "Flow", "published": 1, "latest_version": 1})
    now = "2026-01-01T00:00:00+00:00"
    db.create_run(
        {
            "id": "run_done",
            "workflow_id": "flow",
            "workflow_version": 1,
            "status": "COMPLETED",
            "progress": 1,
            "created_at": now,
            "updated_at": now,
        }
    )
    scheduler = WorkflowScheduler(db, tmp_path / "storage")
    scheduler.execute_run("missing")
    scheduler.execute_run("run_done")
    assert db.get_run("run_done")["status"] == "COMPLETED"


def test_execute_missing_input(tmp_path) -> None:
    """验证输入引用无法解析时任务失败。"""
    db = _db(tmp_path)
    definition = WorkflowDefinition(
        name="missing-input",
        version=1,
        nodes=[
            WorkflowNode(
                id="step",
                node_type="echo",
                inputs={"text": "missing.output"},
            )
        ],
        entry_inputs={"video_uri": "file"},
    )
    db.upsert_workflow({"id": "flow", "name": "Flow", "published": 1, "latest_version": 1})
    db.create_workflow_version("flow", 1, definition.to_dict())
    now = "2026-01-01T00:00:00+00:00"
    db.create_run(
        {
            "id": "run_input",
            "workflow_id": "flow",
            "workflow_version": 1,
            "status": "QUEUED",
            "progress": 0,
            "input_uri": str(tmp_path / "in.txt"),
            "created_at": now,
            "updated_at": now,
        }
    )
    scheduler = WorkflowScheduler(db, tmp_path / "storage")
    scheduler.execute_run("run_input")
    assert db.get_run("run_input")["status"] == "FAILED"


def test_execute_node_failed_response(tmp_path) -> None:
    """验证节点返回 failed 时任务被标记为失败。"""
    db = _db(tmp_path)
    registry.register(
        NodeManifest(
            id="fail-node",
            name="Fail",
            version="1",
            capability="echo",
            repo_dir="nodes",
            command=["python", "-m", "fail"],
        ),
        lambda request: InvokeResponse(status="failed", error="boom"),
    )
    definition = WorkflowDefinition(
        name="fail-flow",
        version=1,
        nodes=[
            WorkflowNode(
                id="step",
                node_type="fail-node",
                inputs={"text": "input.video_uri"},
            )
        ],
        entry_inputs={"video_uri": "file"},
    )
    db.upsert_workflow({"id": "flow", "name": "Flow", "published": 1, "latest_version": 1})
    db.create_workflow_version("flow", 1, definition.to_dict())
    now = "2026-01-01T00:00:00+00:00"
    db.create_run(
        {
            "id": "run_fail_node",
            "workflow_id": "flow",
            "workflow_version": 1,
            "status": "QUEUED",
            "progress": 0,
            "input_uri": str(tmp_path / "in.txt"),
            "created_at": now,
            "updated_at": now,
        }
    )
    scheduler = WorkflowScheduler(db, tmp_path / "storage")
    scheduler.execute_run("run_fail_node")
    assert db.get_run("run_fail_node")["status"] == "FAILED"


def test_scheduler_start_stop_loop(tmp_path) -> None:
    """验证调度线程可重复启动并正常停止。"""
    db = _db(tmp_path)
    scheduler = WorkflowScheduler(db, tmp_path / "storage", interval_seconds=0.05)
    scheduler.start()
    try:
        scheduler.start()
        time.sleep(0.15)
    finally:
        scheduler.stop()
    assert scheduler._thread is None


def test_scheduler_background_executes_queued_run(tmp_path) -> None:
    """验证后台线程会自动执行排队中的任务。"""
    db = _db(tmp_path)
    input_file = tmp_path / "input.txt"
    input_file.write_text("background", encoding="utf-8")
    _register_echo()
    db.upsert_workflow({"id": "flow", "name": "Flow", "published": 1, "latest_version": 1})
    db.create_workflow_version("flow", 1, _echo_definition().to_dict())
    now = "2026-01-01T00:00:00+00:00"
    db.create_run(
        {
            "id": "run_bg",
            "workflow_id": "flow",
            "workflow_version": 1,
            "status": "QUEUED",
            "progress": 0,
            "input_uri": str(input_file),
            "created_at": now,
            "updated_at": now,
        }
    )
    scheduler = WorkflowScheduler(db, tmp_path / "storage", interval_seconds=0.05)
    scheduler.start()
    try:
        deadline = time.monotonic() + 10
        while time.monotonic() < deadline:
            if db.get_run("run_bg")["status"] in {"COMPLETED", "FAILED"}:
                break
            time.sleep(0.1)
    finally:
        scheduler.stop()
    assert db.get_run("run_bg")["status"] == "COMPLETED"


def test_final_artifact_renamed_with_language_tag(tmp_path) -> None:
    """验证最终产物按 上传文件名.语言.时间戳 重命名并登记新 URI。"""
    db = _db(tmp_path)
    input_file = tmp_path / "movie01.mp4"
    input_file.write_text("video", encoding="utf-8")
    _register_echo()
    definition = WorkflowDefinition(
        name="lang-flow",
        version=1,
        nodes=[
            WorkflowNode(
                id="step",
                node_type="echo",
                params={"target_language": "zh-CN"},
                inputs={"file_uri": "input.video_uri"},
            )
        ],
        edges=[],
        entry_inputs={"video_uri": "file"},
        final_outputs={"cn_srt": "step.file_uri"},
    )
    db.upsert_workflow({"id": "flow", "name": "Flow", "published": 1, "latest_version": 1})
    db.create_workflow_version("flow", 1, definition.to_dict())
    now = "2026-01-01T00:00:00+00:00"
    db.create_run(
        {
            "id": "run_1",
            "workflow_id": "flow",
            "workflow_version": 1,
            "status": "QUEUED",
            "progress": 0,
            "input_uri": str(input_file),
            "created_at": now,
            "updated_at": now,
        }
    )
    scheduler = WorkflowScheduler(db, tmp_path / "storage")
    scheduler.execute_run("run_1")
    artifacts = db.list_artifacts("run_1")
    final = next(item for item in artifacts if item["name"] == "cn_srt")
    filename = Path(final["uri"]).name
    # 命名规则：movie01.zh-CN.<14位时间戳>.txt
    assert filename.startswith("movie01.zh-CN.")
    assert filename.endswith(".txt")
    assert Path(final["uri"]).is_file()
    # 原始未重命名文件不应残留。
    step_artifacts = [item for item in artifacts if item["name"] == "step.file_uri"]
    assert not Path(step_artifacts[0]["uri"]).exists()


def test_final_artifact_renamed_fallback_base_and_tag(tmp_path) -> None:
    """验证无上传文件时基础名回退 subtitle，无语言参数时标识回退别名。"""
    db = _db(tmp_path)
    _register_echo()
    definition = WorkflowDefinition(
        name="fallback-flow",
        version=1,
        nodes=[
            # 空输入让 echo 走默认文本路径，避免 input_uri 缺失导致解析失败。
            WorkflowNode(id="step", node_type="echo", inputs={})
        ],
        entry_inputs={"video_uri": "file"},
        final_outputs={"result": "step.file_uri"},
    )
    db.upsert_workflow({"id": "flow", "name": "Flow", "published": 1, "latest_version": 1})
    db.create_workflow_version("flow", 1, definition.to_dict())
    now = "2026-01-01T00:00:00+00:00"
    # 故意不提供 input_uri，验证基础名回退。
    db.create_run(
        {
            "id": "run_1",
            "workflow_id": "flow",
            "workflow_version": 1,
            "status": "QUEUED",
            "progress": 0,
            "input_uri": None,
            "created_at": now,
            "updated_at": now,
        }
    )
    scheduler = WorkflowScheduler(db, tmp_path / "storage")
    scheduler.execute_run("run_1")
    final = next(item for item in db.list_artifacts("run_1") if item["name"] == "result")
    filename = Path(final["uri"]).name
    # 基础名回退 subtitle、标识回退别名 result。
    assert filename.startswith("subtitle.result.")
    assert filename.endswith(".txt")


def test_execute_run_merges_param_overrides(tmp_path) -> None:
    """验证调度执行时把 param_overrides 合并进节点参数。"""
    db = _db(tmp_path)
    input_file = tmp_path / "input.txt"
    input_file.write_text("x", encoding="utf-8")
    captured = {}

    def recording_handler(request):
        captured["params"] = dict(request.params)
        return InvokeResponse(status="completed", outputs={"text": "ok"})

    registry.register(
        NodeManifest(
            id="record-node",
            name="Record",
            version="1",
            capability="echo",
            repo_dir="nodes",
            command=["python", "-m", "record"],
        ),
        recording_handler,
    )
    definition = WorkflowDefinition(
        name="ov-flow",
        version=1,
        nodes=[WorkflowNode(id="step", node_type="record-node", params={"base": 1})],
        entry_inputs={"video_uri": "file"},
    )
    db.upsert_workflow({"id": "flow", "name": "F", "published": 1, "latest_version": 1})
    db.create_workflow_version("flow", 1, definition.to_dict())
    now = "2026-01-01T00:00:00+00:00"
    db.create_run(
        {
            "id": "run_ov",
            "workflow_id": "flow",
            "workflow_version": 1,
            "status": "QUEUED",
            "progress": 0,
            "param_overrides": {"step": {"crop": [0, 0.5, 1, 0.5]}},
            "input_uri": str(input_file),
            "created_at": now,
            "updated_at": now,
        }
    )
    scheduler = WorkflowScheduler(db, tmp_path / "storage")
    scheduler.execute_run("run_ov")
    assert captured["params"] == {"base": 1, "crop": [0, 0.5, 1, 0.5]}


def _two_node_definition() -> WorkflowDefinition:
    """构造 a→b 两节点工作流：b 引用 a 的输出。"""
    return WorkflowDefinition(
        name="two-flow",
        version=1,
        nodes=[
            WorkflowNode(id="a", node_type="x", inputs={"video_uri": "input.video_uri"}),
            WorkflowNode(id="b", node_type="x", inputs={"data_uri": "a.data_uri"}),
        ],
        edges=[WorkflowEdge(from_node="a", to_node="b")],
        entry_inputs={"video_uri": "file"},
        final_outputs={"result": "b.data_uri"},
    )


def test_execute_pause_between_nodes_and_resume(tmp_path, monkeypatch) -> None:
    """验证运行中暂停：节点边界停下保持 PAUSED；续跑时跳过已完成节点。"""
    db = _db(tmp_path)
    db.upsert_workflow({"id": "flow", "name": "Flow", "published": 1, "latest_version": 1})
    db.create_workflow_version("flow", 1, _two_node_definition().to_dict())
    now = "2026-01-01T00:00:00+00:00"
    db.create_run(
        {
            "id": "run_pause",
            "workflow_id": "flow",
            "workflow_version": 1,
            "status": "QUEUED",
            "progress": 0,
            "input_uri": str(tmp_path / "in.txt"),
            "created_at": now,
            "updated_at": now,
        }
    )
    data_file = tmp_path / "data.bin"
    data_file.write_bytes(b"x")

    calls: list[str] = []

    def fake_invoke(node_type, request):
        # registry.invoke 首参是 node_type；用产物目录名（steps/<node_id>）识别节点。
        node_id = Path(request.output_dir).name
        calls.append(node_id)
        # 第一个节点完成后立刻暂停任务，模拟用户在运行中点暂停。
        if node_id == "a":
            db.pause_run("run_pause", now)
        return InvokeResponse(status="completed", outputs={"data_uri": str(data_file)})

    monkeypatch.setattr(registry, "invoke", fake_invoke)
    scheduler = WorkflowScheduler(db, tmp_path / "storage")
    scheduler.execute_run("run_pause")
    assert db.get_run("run_pause")["status"] == "PAUSED"
    assert calls == ["a"]  # 节点 b 未执行。

    # 继续：恢复排队并再次执行，节点 a 已产出结果应被跳过，只执行 b。
    db.resume_run("run_pause", now)
    scheduler.execute_run("run_pause")
    assert db.get_run("run_pause")["status"] == "COMPLETED"
    assert calls == ["a", "b"]
    artifacts = db.list_artifacts("run_pause")
    assert {item["name"] for item in artifacts} == {"a.data_uri", "b.data_uri", "result"}


def test_execute_paused_run_not_run(tmp_path, monkeypatch) -> None:
    """验证非可执行状态（如 RUNNING 之外的值）的任务不会被执行。"""
    db = _db(tmp_path)
    db.upsert_workflow({"id": "flow", "name": "Flow", "published": 1, "latest_version": 1})
    db.create_workflow_version("flow", 1, _two_node_definition().to_dict())
    now = "2026-01-01T00:00:00+00:00"
    db.create_run(
        {
            "id": "run_done",
            "workflow_id": "flow",
            "workflow_version": 1,
            "status": "COMPLETED",
            "progress": 1.0,
            "created_at": now,
            "updated_at": now,
        }
    )
    called = []

    def fake_invoke(node_id, request):
        called.append(node_id)
        return InvokeResponse(status="completed", outputs={})

    monkeypatch.setattr(registry, "invoke", fake_invoke)
    WorkflowScheduler(db, tmp_path / "storage").execute_run("run_done")
    assert called == []


def test_execute_pause_after_last_node_keeps_paused(tmp_path, monkeypatch) -> None:
    """验证全部节点完成但运行中被暂停时保持 PAUSED；续跑补做收尾后完成。"""
    db = _db(tmp_path)
    db.upsert_workflow({"id": "flow", "name": "Flow", "published": 1, "latest_version": 1})
    db.create_workflow_version("flow", 1, _two_node_definition().to_dict())
    now = "2026-01-01T00:00:00+00:00"
    db.create_run(
        {
            "id": "run_tail",
            "workflow_id": "flow",
            "workflow_version": 1,
            "status": "QUEUED",
            "progress": 0,
            "input_uri": str(tmp_path / "in.txt"),
            "created_at": now,
            "updated_at": now,
        }
    )
    data_file = tmp_path / "data.bin"
    data_file.write_bytes(b"x")
    calls: list[str] = []

    def fake_invoke(node_type, request):
        node_id = Path(request.output_dir).name
        calls.append(node_id)
        if node_id == "b":  # 最后一个节点执行时暂停。
            db.pause_run("run_tail", now)
        return InvokeResponse(status="completed", outputs={"data_uri": str(data_file)})

    monkeypatch.setattr(registry, "invoke", fake_invoke)
    scheduler = WorkflowScheduler(db, tmp_path / "storage")
    scheduler.execute_run("run_tail")
    # 全部节点已执行，但收尾前被暂停 → 保持 PAUSED 而不是 COMPLETED。
    assert db.get_run("run_tail")["status"] == "PAUSED"
    assert calls == ["a", "b"]

    # 续跑：节点产物齐备全部跳过，补做收尾后完成。
    db.resume_run("run_tail", now)
    scheduler.execute_run("run_tail")
    assert db.get_run("run_tail")["status"] == "COMPLETED"
    assert calls == ["a", "b"]