rig-pocket/test_server.py

import unittest

import server


def snap(commits=(), queue=(), waiting=False):
    return {
        "rigs": [{"name": "team", "seats": [{"name": "dev-qa@team", "waiting": waiting}]}],
        "queue": list(queue),
        "repos": [{"name": "repo", "commits": [{"hash": h, "subject": f"commit {h}"} for h in commits]}],
    }


def q(qid, state, dest="orch-lead@team"):
    return {"id": qid, "state": state, "dest": dest, "source": "operator-agent@kernel", "summary": f"task {qid}"}


class WatcherTest(unittest.TestCase):
    def test_first_snapshot_only_seeds(self):
        w = server.Watcher()
        self.assertEqual(w.events(snap(commits=["a"], queue=[q("1", "in-progress")])), [])

    def test_new_commits_grouped_per_repo(self):
        w = server.Watcher()
        w.events(snap(commits=["a"]))
        ev = w.events(snap(commits=["a", "b", "c"]))
        self.assertEqual(len(ev), 1)
        self.assertIn("2 ny commits", ev[0][0])

    def test_lead_task_done_notifies_once(self):
        w = server.Watcher()
        w.events(snap(queue=[q("1", "in-progress"), q("2", "in-progress", "dev-build@team")]))
        ev = w.events(snap(queue=[q("1", "done"), q("2", "done", "dev-build@team")]))
        self.assertEqual([e[0] for e in ev], ["Opgave done: orch-lead@team"])
        self.assertEqual(w.events(snap(queue=[q("1", "done")])), [])

    def test_new_request_to_human_notifies(self):
        w = server.Watcher()
        w.events(snap())
        ev = w.events(snap(queue=[q("9", "in-progress", "operator-agent@kernel")]))
        self.assertEqual(ev[0][0], "Teamet spørger dig")

    def test_prompt_must_survive_two_polls(self):
        w = server.Watcher()
        self.assertEqual(w.events(snap(waiting=True)), [])
        ev = w.events(snap(waiting=True))
        self.assertEqual(ev[0][0], "dev-qa@team venter på svar")
        self.assertEqual(w.events(snap(waiting=True)), [])
        w.events(snap(waiting=False))
        w.events(snap(waiting=True))
        self.assertEqual(len(w.events(snap(waiting=True))), 1)

    def test_gate_notifies_once(self):
        w = server.Watcher()
        w.events(snap())
        gate = {**q("7", "blocked", "operator-human@kernel"), "gate": True}
        self.assertEqual(w.events(snap(queue=[gate]))[-1][0], "Beslutning til dig")
        self.assertEqual(w.events(snap(queue=[gate])), [])


PERMISSION = """[rig capture: operator-agent@kernel: 23 earlier lines not shown]
● Bash(rig queue show qitem-1 --json)
  ⎿  Waiting…

────────────────────────────────────────────────────────────────────────────────
 Bash command
 Tip: auto mode handles these prompts for you — choose "switch to auto mode"
 below
 Check whether the Chromium libraries are installed
╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌
 │ dpkg-query -W -f='${Package} ${Status}\\n' libatk1.0-0t64
╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌
 This command requires approval

 Do you want to proceed?
 ❯ 1. Yes
   2. Yes, and don’t ask again for: dpkg-query *
   3. Yes, and switch to auto mode · auto mode handles these prompts for you
   4. No

 Esc to cancel · Tab to amend
"""

QUESTION = """  - Install location: ~/rigs/workshop.

● Once it's running: I'll give the lead the whole task as a durable work item
  (spec file, branch, local commits only, no push). When it's done, I'll report
  back the commits and test results.

● Should I launch workshop like this (all Claude, prompts off, in
  /home/rig/work/wpp)?

✻ Churned for 1m 24s · done 4:26 PM

──────────────────────────────────────────────────────── operator-agent@kernel ─
❯ yes, launch workshop all-claude
────────────────────────────────────────────────────────────────────────────────
  ⏵⏵ accept edits on (shift+tab to cycle)
"""

WORKING = """● Bash(cargo test --locked)
  ⎿  Running…

✻ Crunching… (1m 7s · ↓ 3.3k tokens · thinking)

──────────────────────────────────────────────────────── dev-build@team ─
❯
────────────────────────────────────────────────────────────────────────────────
"""


class ParseAskTest(unittest.TestCase):
    def test_permission_prompt(self):
        a = server.parse_ask(PERMISSION, idle=False)
        self.assertEqual(a["kind"], "menu")
        self.assertEqual(a["title"], "Do you want to proceed?")
        self.assertEqual([o["key"] for o in a["options"]], ["1", "2", "3", "4"])
        self.assertEqual(a["options"][1]["label"], "Yes, and don’t ask again for: dpkg-query *")

    def test_question_when_idle(self):
        a = server.parse_ask(QUESTION, idle=True)
        self.assertEqual(a["kind"], "question")
        self.assertTrue(a["text"].startswith("Should I launch workshop"))
        self.assertTrue(a["text"].endswith("?"))

    def test_question_needs_idle(self):
        self.assertIsNone(server.parse_ask(QUESTION, idle=False))

    def test_working_is_not_a_question(self):
        self.assertIsNone(server.parse_ask(WORKING, idle=True))
        self.assertIsNone(server.parse_ask(WORKING, idle=False))

    def test_numbered_list_without_cursor_is_not_a_menu(self):
        self.assertIsNone(server.parse_ask("● Plan:\n  1. engine\n  2. rooms\n", idle=False))

    def test_stalled_turn(self):
        self.assertEqual(server.stalled("· Ionizing… (34m 31s · ↓ 439 tokens · thinking some more)"), 35)
        self.assertIsNone(server.stalled("✶ Levitating… (10m 40s · ↓ 30.1k tokens)"))  # short and busy
        self.assertIsNone(server.stalled("✻ Doing… (16m 3s · ↓ 3.4k tokens)"))  # long but producing
        self.assertEqual(server.stalled("* Thinking… (1h 2m 0s · ↓ 900 tokens)"), 62)
        self.assertIsNone(server.stalled("✻ Churned for 1m 24s · done 4:26 PM"))

    def test_stall_notifies_once(self):
        w = server.Watcher()
        s = snap()
        s["rigs"][0]["seats"][0]["stalled"] = 34
        self.assertEqual(w.events(s)[-1][0], "dev-qa@team ser ud til at hænge")
        self.assertEqual(w.events(s), [])

    def test_memory_warning_once(self):
        w = server.Watcher()
        s = {**snap(), "memory": {"used": 95, "limit": 100, "fraction": 0.95}}
        self.assertTrue(w.events(s)[-1][0].startswith("Containeren"))
        self.assertEqual(w.events(s), [])

    def test_ask_text_for_notifications(self):
        self.assertIn("2. Yes, and", server.ask_text(server.parse_ask(PERMISSION, idle=False)))

    def test_seat_names(self):
        self.assertTrue(server.SEAT.match("dev-qa@ubl-workshop"))
        self.assertFalse(server.SEAT.match("x; rm -rf /"))
        self.assertFalse(server.SEAT.match("../etc@x/y"))


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