ai-agent/backend/test_workers.py

"""Workers: pairing-string parsing, the store, pairing against a real
(loopback) worker, probing. ``python -m unittest test_workers`` from
``backend/`` with fastapi + uvicorn installed (the service image has both)."""

import pathlib
import tempfile
import unittest

from fastapi import HTTPException

import workers
from _fakeworker import FakeWorker


class ParsePairing(unittest.TestCase):
    def test_full(self):
        p = workers.parse_pairing(" 100.80.162.92:8790/k7qmx4pj2r/g@orus.insure ")
        self.assertEqual(p, {"host": "100.80.162.92", "port": 8790,
                             "code": "K7QMX4PJ2R", "login": "g@orus.insure"})

    def test_no_login_default_port_and_url_prefix(self):
        p = workers.parse_pairing("http://mac.tail.ts.net/ABCDEF2345")
        self.assertEqual((p["host"], p["port"], p["login"]),
                         ("mac.tail.ts.net", 8790, None))

    def test_rejects_garbage(self):
        for bad in ("", "100.1.2.3:8790", "100.1.2.3:99999/ABCDEFGH",
                    "host:port/ABCDEFGH", "a b:1/ABCDEFGH", "1.2.3.4:1/ab"):
            with self.assertRaises(HTTPException, msg=bad) as cm:
                workers.parse_pairing(bad)
            self.assertEqual(cm.exception.detail["kind"], "pairing")


class PairAgainstAWorker(unittest.TestCase):
    def setUp(self):
        self.tmp = pathlib.Path(tempfile.mkdtemp())
        self.fw = FakeWorker(self.tmp / "projects", self.tmp / "state")
        self.store = workers.WorkerStore(str(self.tmp / "workers.json"))

    def tearDown(self):
        self.fw.stop()

    def _pair(self, raw, account=None, resolve=lambda login: None):
        return workers.pair(self.store, raw, account=account, name=None,
                            hub_name="test-hub", resolve_login=resolve,
                            valid_account=lambda a: a if a in (
                                "personal", "work") else None)

    def test_login_resolves_the_account_and_pairs(self):
        w = self._pair(self.fw.pairing_string(),
                       resolve=lambda login: "work"
                       if login == "g@orus.example" else None)
        self.assertEqual(w["account"], "work")
        self.assertEqual(w["email"], "g@orus.example")
        self.assertEqual(w["cwd"], "/Users/g/repo")
        self.assertTrue(w["token"])
        self.assertTrue(self.store.online(w["id"]))
        # persisted, 0600, token never in the public view
        self.assertEqual((self.tmp / "workers.json").stat().st_mode & 0o777,
                         0o600)
        self.assertNotIn("token", self.store.public(w))
        # the code is spent
        with self.assertRaises(HTTPException) as cm:
            self._pair(f"127.0.0.1:{self.fw.port}/ABCDEFGHJK/work")
        self.assertEqual(cm.exception.status_code, 409)

    def test_unmatched_login_asks_for_an_account_without_spending_the_code(self):
        raw = self.fw.pairing_string("someone@else.example")
        with self.assertRaises(HTTPException) as cm:
            self._pair(raw)
        self.assertEqual(cm.exception.detail["kind"], "account")
        w = self._pair(raw, account="work")      # same code still good
        self.assertEqual(w["account"], "work")

    def test_login_segment_may_be_an_account_id(self):
        self.assertEqual(self._pair(self.fw.pairing_string("work"))["account"],
                         "work")

    def test_unreachable_worker_is_worker_offline(self):
        self.fw.stop()
        with self.assertRaises(HTTPException) as cm:
            self._pair(f"127.0.0.1:{self.fw.port}/ABCDEFGHJK", account="work")
        self.assertEqual(cm.exception.detail["kind"], "worker_offline")

    def test_probe_tracks_online_offline_and_running(self):
        w = self._pair(self.fw.pairing_string("work"))
        self.fw.live = {"s1", "s2"}
        workers.probe(self.store, w)
        self.assertEqual(self.store.status(w["id"])["running"], 2)
        # a rotated token (re-paired elsewhere) reads as unauthorized
        self.store.patch(w["id"], {"token": "stale"})
        workers.probe(self.store, self.store.get(w["id"]))
        self.assertEqual(self.store.status(w["id"])["status"], "unauthorized")
        self.fw.stop()
        workers.probe(self.store, self.store.get(w["id"]))
        self.assertEqual(self.store.status(w["id"])["status"], "offline")

    def test_calls_carry_the_bearer_and_map_errors(self):
        w = self._pair(self.fw.pairing_string("work"))
        out = workers.sidecar_request(w["url"], w["token"], "/sessions",
                                      method="GET", label="worker mac")
        self.assertEqual(out["count"], 0)
        with self.assertRaises(HTTPException) as cm:
            workers.sidecar_request(w["url"], "nope", "/sessions",
                                    method="GET", label="worker mac")
        self.assertEqual((cm.exception.status_code,
                          cm.exception.detail["kind"]), (502, "worker_auth"))

    def test_unpair_drops_the_record_and_the_workers_bearer(self):
        w = self._pair(self.fw.pairing_string("work"))
        self.assertTrue(workers.unpair(self.store, w["id"]))
        self.assertIsNone(self.store.get(w["id"]))
        import pairing
        self.assertEqual(pairing.token(), "")

    def test_for_account_prefers_online_workers(self):
        w = self._pair(self.fw.pairing_string("work"))
        self.assertEqual(self.store.for_account("work")["id"], w["id"])
        self.assertIsNone(self.store.for_account("personal"))
        self.store.set_status(w["id"], status="offline")
        self.assertIsNone(self.store.for_account("work"))


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