local-rig/tests/test_local_rig.py

"""local_rig 测试：schema 解析 + 队列持久化 + daemon API。"""
from __future__ import annotations

import json
import os
import socket
import sqlite3
import tempfile
import threading
import time
import unittest
import urllib.request
import ssl
from pathlib import Path
from unittest.mock import patch

# ── schema 测试 ──

import sys
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from local_rig.rig_schema import parse_rig, Rig, RigSchemaError, Member, Pod, Edge


def _write_rig(tmp: str, content: str) -> str:
    p = os.path.join(tmp, "rig.yaml")
    with open(p, "w") as f:
        f.write(content)
    return p


MINIMAL_RIG = """
version: "0.1"
name: test-rig
pods:
  - id: dev
    members:
      - id: lead
        runtime: terminal
        command: "echo hi"
"""

FULL_RIG = """
version: "0.1"
name: full-rig
pods:
  - id: dev
    label: Development
    members:
      - id: lead
        runtime: hermes-cli
        target: hermes
      - id: qa
        runtime: terminal
        command: "python3 qa.py"
        cwd: ./qa
    edges:
      - kind: delegates_to
        from: lead
        to: qa
  - id: rev
    members:
      - id: r1
        runtime: terminal
        command: "echo review"
    edges:
      - kind: can_observe
        from: r1
        to: dev.lead
"""


class TestRigSchema(unittest.TestCase):

    def setUp(self):
        self.tmp = tempfile.mkdtemp(prefix="hermes-rig-test-")

    def tearDown(self):
        import shutil
        shutil.rmtree(self.tmp, ignore_errors=True)

    def test_minimal_rig(self):
        p = _write_rig(self.tmp, MINIMAL_RIG)
        rig = parse_rig(p)
        self.assertEqual(rig.name, "test-rig")
        self.assertEqual(len(rig.pods), 1)
        self.assertEqual(rig.pods[0].id, "dev")
        members = rig.all_members()
        self.assertIn("dev.lead", members)
        self.assertEqual(members["dev.lead"].runtime, "terminal")

    def test_full_rig(self):
        p = _write_rig(self.tmp, FULL_RIG)
        rig = parse_rig(p)
        members = rig.all_members()
        self.assertEqual(len(members), 3)
        self.assertEqual(members["dev.lead"].runtime, "hermes-cli")
        self.assertEqual(members["dev.lead"].target, "hermes")
        self.assertEqual(members["dev.qa"].cwd, "./qa")
        # edges
        self.assertEqual(len(rig.pods[0].edges), 1)
        e = rig.pods[0].edges[0]
        self.assertEqual(e.kind, "delegates_to")
        self.assertEqual(rig.resolve_edge(e), ("dev.lead", "dev.qa"))

    def test_missing_name(self):
        p = _write_rig(self.tmp, "pods: []\n")
        with self.assertRaises(RigSchemaError):
            parse_rig(p)

    def test_empty_pods(self):
        p = _write_rig(self.tmp, 'name: x\npods: []\n')
        with self.assertRaises(RigSchemaError):
            parse_rig(p)

    def test_bad_runtime(self):
        content = MINIMAL_RIG.replace("terminal", "warp-drive")
        p = _write_rig(self.tmp, content)
        with self.assertRaises(RigSchemaError):
            parse_rig(p)

    def test_terminal_requires_command(self):
        content = """
name: t
pods:
  - id: dev
    members:
      - id: x
        runtime: terminal
"""
        p = _write_rig(self.tmp, content)
        with self.assertRaises(RigSchemaError):
            parse_rig(p)

    def test_hermes_requires_target(self):
        content = """
name: t
pods:
  - id: dev
    members:
      - id: x
        runtime: hermes-cli
"""
        p = _write_rig(self.tmp, content)
        with self.assertRaises(RigSchemaError):
            parse_rig(p)

    def test_duplicate_pod_id(self):
        content = """
name: t
pods:
  - id: dev
    members:
      - id: a
        runtime: terminal
        command: "echo"
  - id: dev
    members:
      - id: b
        runtime: terminal
        command: "echo"
"""
        p = _write_rig(self.tmp, content)
        with self.assertRaises(RigSchemaError):
            parse_rig(p)

    def test_bad_edge_ref(self):
        content = """
name: t
pods:
  - id: dev
    members:
      - id: a
        runtime: terminal
        command: "echo"
    edges:
      - kind: delegates_to
        from: a
        to: ghost
"""
        p = _write_rig(self.tmp, content)
        with self.assertRaises(RigSchemaError):
            parse_rig(p)

    def test_bad_id_chars(self):
        content = """
name: t
pods:
  - id: Dev
    members:
      - id: a
        runtime: terminal
        command: "echo"
"""
        p = _write_rig(self.tmp, content)
        with self.assertRaises(RigSchemaError):
            parse_rig(p)

    def test_cross_pod_edge(self):
        p = _write_rig(self.tmp, FULL_RIG)
        rig = parse_rig(p)
        e = rig.pods[1].edges[0]  # rev.r1 → dev.lead
        self.assertEqual(rig.resolve_edge(e), ("rev.r1", "dev.lead"))


# ── 队列测试 ──

from local_rig.queue import TaskQueue


class TestTaskQueue(unittest.TestCase):

    def setUp(self):
        self.tmp = tempfile.mkdtemp(prefix="hermes-rig-queue-")
        self.db = os.path.join(self.tmp, "test.sqlite3")
        self.q = TaskQueue(self.db)

    def tearDown(self):
        self.q.close()
        import shutil
        shutil.rmtree(self.tmp, ignore_errors=True)

    def test_enqueue_and_get(self):
        t = self.q.enqueue("dev.lead", "修 bug")
        self.assertEqual(t.status, "pending")
        got = self.q.get(t.id)
        self.assertIsNotNone(got)
        self.assertEqual(got.prompt, "修 bug")
        self.assertEqual(got.assignee, "dev.lead")

    def test_idempotent_enqueue(self):
        t1 = self.q.enqueue("dev.lead", "修 bug", task_id="fixed-id")
        t2 = self.q.enqueue("dev.lead", "不同的 prompt", task_id="fixed-id")
        self.assertEqual(t1.id, t2.id)
        self.assertEqual(t2.prompt, "修 bug")  # 不覆盖
        self.assertEqual(self.q.count(), 1)

    def test_claim_atomic(self):
        self.q.enqueue("dev.lead", "task1")
        t = self.q.claim(assignee="dev.lead")
        self.assertIsNotNone(t)
        self.assertEqual(t.status, "running")
        # 再领 → None
        self.assertIsNone(self.q.claim(assignee="dev.lead"))

    def test_claim_wrong_assignee(self):
        self.q.enqueue("dev.lead", "task1")
        self.assertIsNone(self.q.claim(assignee="rev.r1"))

    def test_complete(self):
        t = self.q.enqueue("dev.lead", "task1")
        self.q.claim(assignee="dev.lead")
        self.q.complete(t.id, result="done!")
        got = self.q.get(t.id)
        self.assertEqual(got.status, "done")
        self.assertEqual(got.result, "done!")

    def test_fail(self):
        t = self.q.enqueue("dev.lead", "task1")
        self.q.claim(assignee="dev.lead")
        self.q.fail(t.id, error="boom")
        got = self.q.get(t.id)
        self.assertEqual(got.status, "failed")
        self.assertEqual(got.error, "boom")

    def test_timeout_sweep(self):
        t = self.q.enqueue("dev.lead", "task1", deadline_s=0.01)
        self.q.claim(assignee="dev.lead")
        time.sleep(0.05)
        timed_out = self.q.sweep_timeouts()
        self.assertIn(t.id, timed_out)
        got = self.q.get(t.id)
        self.assertEqual(got.status, "timeout")

    def test_persistence_across_restart(self):
        t = self.q.enqueue("dev.lead", "持久化测试")
        self.q.close()
        # 重开
        q2 = TaskQueue(self.db)
        got = q2.get(t.id)
        self.assertIsNotNone(got)
        self.assertEqual(got.prompt, "持久化测试")
        self.assertEqual(got.status, "pending")
        q2.close()

    def test_list_filter(self):
        self.q.enqueue("dev.lead", "a")
        self.q.enqueue("dev.qa", "b")
        self.q.enqueue("dev.lead", "c")
        self.assertEqual(len(self.q.list(assignee="dev.lead")), 2)
        self.assertEqual(len(self.q.list(status="pending")), 3)

    def test_meta(self):
        self.q.set_meta("version", "0.1")
        self.assertEqual(self.q.get_meta("version"), "0.1")
        self.assertIsNone(self.q.get_meta("nonexistent"))


# ── daemon API 测试 ──

from local_rig.daemon import RigDaemon


class TestRigDaemon(unittest.TestCase):

    def setUp(self):
        self.tmp = tempfile.mkdtemp(prefix="hermes-rig-daemon-")
        self.rig_path = _write_rig(self.tmp, MINIMAL_RIG)
        self.db_path = os.path.join(self.tmp, "rig.sqlite3")
        # 动态端口（避免冲突）
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.bind(("127.0.0.1", 0))
        self.port = s.getsockname()[1]
        s.close()
        self.d = RigDaemon(self.rig_path, db_path=self.db_path, port=self.port)
        self.d.start(background=True)
        self.ctx = ssl.create_default_context()
        self.ctx.check_hostname = False
        self.ctx.verify_mode = ssl.CERT_NONE
        time.sleep(0.2)  # 等 daemon 就绪

    def tearDown(self):
        self.d.stop()
        import shutil
        shutil.rmtree(self.tmp, ignore_errors=True)

    def _get(self, path):
        req = urllib.request.Request(f"http://127.0.0.1:{self.port}{path}")
        with urllib.request.urlopen(req, context=self.ctx, timeout=5) as resp:
            return json.loads(resp.read())

    def _post(self, path, data=None):
        body = json.dumps(data or {}).encode()
        req = urllib.request.Request(
            f"http://127.0.0.1:{self.port}{path}", data=body, method="POST",
            headers={"Content-Type": "application/json"})
        with urllib.request.urlopen(req, context=self.ctx, timeout=5) as resp:
            return json.loads(resp.read())

    def test_health(self):
        r = self._get("/health")
        self.assertEqual(r["status"], "ok")
        self.assertEqual(r["rig"], "test-rig")

    def test_rig_definition(self):
        r = self._get("/rig")
        self.assertEqual(r["name"], "test-rig")
        self.assertIn("pods", r)

    def test_status(self):
        r = self._get("/status")
        self.assertIn("dev.lead", r)

    def test_enqueue_and_list(self):
        r = self._post("/tasks", {"assignee": "dev.lead", "prompt": "测试任务"})
        self.assertEqual(r["status"], "pending")
        tid = r["id"]
        tasks = self._get("/tasks")
        self.assertEqual(len(tasks), 1)
        self.assertEqual(tasks[0]["prompt"], "测试任务")

    def test_task_lifecycle(self):
        r = self._post("/tasks", {"assignee": "dev.lead", "prompt": "lifecycle"})
        tid = r["id"]
        # claim
        self._post(f"/tasks/{tid}/claim")
        t = self._get(f"/tasks/{tid}")
        self.assertEqual(t["status"], "running")
        # done
        self._post(f"/tasks/{tid}/done", {"result": "完成"})
        t = self._get(f"/tasks/{tid}")
        self.assertEqual(t["status"], "done")
        self.assertEqual(t["result"], "完成")

    def test_task_fail(self):
        r = self._post("/tasks", {"assignee": "dev.lead", "prompt": "will-fail"})
        tid = r["id"]
        self._post(f"/tasks/{tid}/claim")
        self._post(f"/tasks/{tid}/fail", {"error": "崩溃了"})
        t = self._get(f"/tasks/{tid}")
        self.assertEqual(t["status"], "failed")

    def test_missing_fields(self):
        import urllib.error
        try:
            self._post("/tasks", {"assignee": "dev.lead"})  # 缺 prompt
            self.fail("应该 400")
        except urllib.error.HTTPError as e:
            self.assertEqual(e.code, 400)

    def test_404_unknown_task(self):
        import urllib.error
        try:
            self._get("/tasks/nonexistent")
            self.fail("应该 404")
        except urllib.error.HTTPError as e:
            self.assertEqual(e.code, 404)


if __name__ == "__main__":
    unittest.main()


class TestMemberResolution(unittest.TestCase):
    """实测缺口回归：assignee 传裸 member id 时曾**静默不派发**（任务永远 pending）。"""

    def _daemon(self, runtime="terminal"):
        from local_rig.daemon import RigDaemon
        from pathlib import Path
        cfg = Path(tempfile.mkdtemp()) / "rig.yaml"
        cfg.write_text(
            "version: \"0.1\"\nname: t\npods:\n  - id: dev\n    label: D\n    members:\n"
            "      - id: hermes\n        runtime: %s\n        target: hermes\n"
            "        command: \"echo hi\"\n    edges: []\n" % runtime, encoding="utf-8")
        return RigDaemon(cfg, db_path=str(Path(tempfile.mkdtemp()) / "q.db"))

    def test_bare_member_id_resolves_to_pod_member(self):
        d = self._daemon()
        self.assertEqual(d.resolve_member("hermes"), "dev.hermes")
        self.assertEqual(d.resolve_member("dev.hermes"), "dev.hermes")

    def test_unknown_member_raises_instead_of_silent_pending(self):
        d = self._daemon()
        with self.assertRaises(KeyError):
            d.enqueue("nope", "x")

    def test_hermes_cli_dispatch_actually_fires(self):
        """端到端最小闭环：入队 → 真子进程 → done（曾经停在 pending 且无任何报错）。"""
        import os, stat, time
        d = self._daemon(runtime="hermes-cli")
        fake = os.path.join(tempfile.mkdtemp(), "fake-agent")
        with open(fake, "w") as f:
            f.write('#!/bin/sh\necho "AGENT-OK: $*"\n')
        os.chmod(fake, stat.S_IRWXU)
        d.hermes_rt.hermes_bin = fake
        task = d.enqueue("hermes", "说一句你好", task_id="gap-1", idem_key="idem-gap")
        row = None
        for _ in range(20):
            time.sleep(0.5)
            row = d.queue.get(task.id)
            if row and row.status in ("done", "failed"):
                break
        self.assertIsNotNone(row)
        self.assertEqual(row.status, "done", "裸 member id 的任务应被真正执行，而不是停在 pending")
        self.assertIn("AGENT-OK", row.result or "")

    def test_idempotency_survives_real_dispatch_path(self):
        import os, stat
        d = self._daemon(runtime="hermes-cli")
        fake = os.path.join(tempfile.mkdtemp(), "fa"); open(fake, "w").write('#!/bin/sh\necho ok\n')
        os.chmod(fake, stat.S_IRWXU); d.hermes_rt.hermes_bin = fake
        ids = {d.enqueue("hermes", "同一件事", idem_key="same-key").id for _ in range(3)}
        self.assertEqual(len(ids), 1, "同 idem_key 三次 ⇒ 只应有一个任务")