didericis
b1df380ae1
Consolidated proposals are keyed by the orchestrator-assigned bottle_id, so the supervise TUI was rendering a hex id (e.g. 3601cbe883c2786d) as the bottle name — and the resume hint printed `./cli.py resume <id>`, which resume can't take (it wants the human identity/slug). `supervise_pending` now tags each dict with `bottle_label` — the human slug resolved from registry metadata, falling back to the id when the bottle is gone. `QueuedProposal` carries the label; every display site (list rows, detail view, status lines, resume hint) shows it, while approve/reject still key on `proposal.bottle_slug` (the id). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UoEZHDjv84ChoZbozQERhJ
307 lines
13 KiB
Python
307 lines
13 KiB
Python
"""Unit tests for the Orchestrator launch lifecycle (PRD 0070)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import secrets
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
from bot_bottle.orchestrator.broker import LaunchBroker, LaunchRequest, StubBroker
|
|
from bot_bottle.orchestrator.registry import RegistryStore
|
|
from bot_bottle.orchestrator.service import Orchestrator
|
|
from bot_bottle.orchestrator.gateway import Gateway
|
|
from bot_bottle.store_manager import StoreManager
|
|
from bot_bottle.supervise import (
|
|
Proposal,
|
|
STATUS_APPROVED,
|
|
TOOL_EGRESS_ALLOW,
|
|
read_response,
|
|
sha256_hex,
|
|
write_proposal,
|
|
)
|
|
|
|
|
|
class _FailingBroker(LaunchBroker):
|
|
"""Verifies the token like any broker, then fails the launch — to
|
|
exercise the orchestrator's registry rollback."""
|
|
|
|
def _launch(self, req: LaunchRequest) -> None:
|
|
raise RuntimeError("launch failed")
|
|
|
|
def _teardown(self, req: LaunchRequest) -> None:
|
|
pass
|
|
|
|
|
|
class _FakeGateway(Gateway):
|
|
"""In-memory gateway for wiring tests."""
|
|
|
|
def __init__(self) -> None:
|
|
self.name = "fake-gateway"
|
|
self.ensured = 0
|
|
self.built = 0
|
|
self._running = False
|
|
|
|
def ensure_built(self) -> None:
|
|
self.built += 1
|
|
|
|
def ensure_running(self) -> None:
|
|
self.ensured += 1
|
|
self._running = True
|
|
|
|
def is_running(self) -> bool:
|
|
return self._running
|
|
|
|
def stop(self) -> None:
|
|
self._running = False
|
|
|
|
|
|
class TestOrchestrator(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
self._tmp = tempfile.TemporaryDirectory()
|
|
self.secret = secrets.token_bytes(16)
|
|
self.store = RegistryStore(Path(self._tmp.name) / "r.db")
|
|
self.store.migrate()
|
|
self.broker = StubBroker(self.secret)
|
|
self.orch = Orchestrator(self.store, self.broker, self.secret)
|
|
|
|
def tearDown(self) -> None:
|
|
self._tmp.cleanup()
|
|
|
|
def test_launch_registers_and_brokers_signed_request(self) -> None:
|
|
rec = self.orch.launch_bottle("10.243.0.1", image_ref="sha256:abc", slot=2)
|
|
self.assertIsNotNone(self.store.get(rec.bottle_id))
|
|
self.assertEqual(1, len(self.broker.launched))
|
|
req = self.broker.launched[0]
|
|
self.assertEqual("launch", req.op)
|
|
self.assertEqual(rec.bottle_id, req.bottle_id)
|
|
self.assertEqual("10.243.0.1", req.source_ip)
|
|
self.assertEqual("sha256:abc", req.image_ref)
|
|
self.assertEqual(2, req.slot)
|
|
|
|
def test_launch_then_attribute(self) -> None:
|
|
rec = self.orch.launch_bottle("10.243.0.3")
|
|
got = self.orch.attribute("10.243.0.3", rec.identity_token)
|
|
assert got is not None
|
|
self.assertEqual(rec.bottle_id, got.bottle_id)
|
|
self.assertIsNone(self.orch.attribute("10.243.0.3", "wrong-token"))
|
|
|
|
def test_teardown_brokers_and_deregisters(self) -> None:
|
|
rec = self.orch.launch_bottle("10.243.0.1")
|
|
self.assertTrue(self.orch.teardown_bottle(rec.bottle_id))
|
|
self.assertIsNone(self.store.get(rec.bottle_id))
|
|
self.assertEqual(["teardown"], [r.op for r in self.broker.torn_down])
|
|
|
|
def test_teardown_unknown_is_false(self) -> None:
|
|
self.assertFalse(self.orch.teardown_bottle("ghost"))
|
|
|
|
def test_launch_with_policy_is_resolvable(self) -> None:
|
|
rec = self.orch.launch_bottle("10.243.0.1", policy='{"routes":[]}')
|
|
got = self.orch.attribute("10.243.0.1", rec.identity_token)
|
|
assert got is not None
|
|
self.assertEqual('{"routes":[]}', got.policy)
|
|
|
|
def test_tokens_held_in_memory_and_cleared_on_teardown(self) -> None:
|
|
rec = self.orch.launch_bottle("10.243.0.5", tokens={"EGRESS_TOKEN_0": "sk"})
|
|
self.assertEqual({"EGRESS_TOKEN_0": "sk"}, self.orch.tokens_for(rec.bottle_id))
|
|
# not persisted in the registry (redacted record has no tokens field)
|
|
self.assertNotIn("tokens", rec.redacted())
|
|
self.orch.teardown_bottle(rec.bottle_id)
|
|
self.assertEqual({}, self.orch.tokens_for(rec.bottle_id))
|
|
|
|
def test_tokens_default_empty(self) -> None:
|
|
rec = self.orch.launch_bottle("10.243.0.6")
|
|
self.assertEqual({}, self.orch.tokens_for(rec.bottle_id))
|
|
|
|
def test_set_policy_live_reload(self) -> None:
|
|
rec = self.orch.launch_bottle("10.243.0.3")
|
|
self.assertTrue(self.orch.set_policy(rec.bottle_id, '{"x":1}'))
|
|
got = self.orch.attribute("10.243.0.3", rec.identity_token)
|
|
assert got is not None
|
|
self.assertEqual('{"x":1}', got.policy)
|
|
|
|
def test_set_policy_unknown_is_false(self) -> None:
|
|
self.assertFalse(self.orch.set_policy("ghost", "{}"))
|
|
|
|
def test_resolve_requires_matching_token(self) -> None:
|
|
# Mandatory (source_ip, token) pair — no source-IP-only fallback.
|
|
rec = self.orch.launch_bottle("10.243.0.1", policy="P")
|
|
self.assertIsNone(self.orch.resolve("10.243.0.1", "")) # empty token denies
|
|
self.assertIsNone(self.orch.resolve("10.243.0.1", "wrong")) # mismatch denies
|
|
got = self.orch.resolve("10.243.0.1", rec.identity_token) # exact pair
|
|
assert got is not None
|
|
self.assertEqual(rec.bottle_id, got.bottle_id)
|
|
self.assertEqual("P", got.policy)
|
|
|
|
def test_resolve_with_token_stays_strict(self) -> None:
|
|
rec = self.orch.launch_bottle("10.243.0.3")
|
|
self.assertIsNotNone(self.orch.resolve("10.243.0.3", rec.identity_token))
|
|
self.assertIsNone(self.orch.resolve("10.243.0.3", "wrong-token"))
|
|
|
|
def test_launch_rolls_back_registry_on_broker_failure(self) -> None:
|
|
orch = Orchestrator(self.store, _FailingBroker(self.secret), self.secret)
|
|
with self.assertRaises(RuntimeError):
|
|
orch.launch_bottle("10.243.0.9")
|
|
self.assertEqual([], self.store.all()) # no orphan
|
|
|
|
def test_gateway_unconfigured_by_default(self) -> None:
|
|
self.assertEqual({"configured": False}, self.orch.gateway_status())
|
|
self.orch.ensure_gateway() # no-op, must not raise
|
|
|
|
def test_gateway_wired_and_ensured(self) -> None:
|
|
sc = _FakeGateway()
|
|
orch = Orchestrator(self.store, self.broker, self.secret, gateway=sc)
|
|
self.assertEqual(
|
|
{"configured": True, "name": "fake-gateway", "running": False},
|
|
orch.gateway_status(),
|
|
)
|
|
orch.ensure_gateway()
|
|
self.assertEqual(1, sc.built) # ensure_gateway builds first,
|
|
self.assertEqual(1, sc.ensured) # then runs
|
|
self.assertEqual(
|
|
{"configured": True, "name": "fake-gateway", "running": True},
|
|
orch.gateway_status(),
|
|
)
|
|
|
|
|
|
class TestOrchestratorSupervise(unittest.TestCase):
|
|
"""Operator-approval flow: the orchestrator applies the decision
|
|
server-side against the single DB (queue + policy + audit)."""
|
|
|
|
def setUp(self) -> None:
|
|
self._tmp = tempfile.TemporaryDirectory()
|
|
root = Path(self._tmp.name)
|
|
db = root / "db" / "bot-bottle.db"
|
|
db.parent.mkdir(parents=True)
|
|
# One DB for registry + supervise queue + audit (as in the VM).
|
|
self._env = patch.dict("os.environ", {
|
|
"BOT_BOTTLE_ROOT": str(root),
|
|
"SUPERVISE_DB_PATH": str(db),
|
|
})
|
|
self._env.start()
|
|
self.store = RegistryStore(db)
|
|
self.store.migrate()
|
|
StoreManager(db).migrate()
|
|
secret = secrets.token_bytes(16)
|
|
self.orch = Orchestrator(self.store, StubBroker(secret), secret)
|
|
|
|
def tearDown(self) -> None:
|
|
self._env.stop()
|
|
self._tmp.cleanup()
|
|
|
|
def _register(self, slug: str, policy: str) -> str:
|
|
rec = self.store.register(
|
|
"10.243.0.1", metadata=json.dumps({"slug": slug}), policy=policy)
|
|
return rec.bottle_id
|
|
|
|
def _queue(self, slug: str, proposed: str) -> str:
|
|
p = Proposal.new(
|
|
bottle_slug=slug, tool=TOOL_EGRESS_ALLOW, proposed_file=proposed,
|
|
justification="need it", current_file_hash=sha256_hex(proposed))
|
|
write_proposal(p)
|
|
return p.id
|
|
|
|
def test_pending_lists_queued_proposal(self) -> None:
|
|
self._register("demo", "routes: []\n")
|
|
pid = self._queue("demo", "routes:\n - host: google.com\n")
|
|
pending = self.orch.supervise_pending()
|
|
self.assertEqual(1, len(pending))
|
|
self.assertEqual(pid, pending[0]["id"])
|
|
self.assertEqual("demo", pending[0]["bottle_slug"])
|
|
|
|
def test_approve_applies_policy_writes_response_and_clears_pending(self) -> None:
|
|
bottle_id = self._register("demo", "routes:\n - host: existing.com\n")
|
|
new_routes = "routes:\n - host: google.com\n"
|
|
pid = self._queue("demo", new_routes)
|
|
ok, err = self.orch.supervise_respond(
|
|
pid, bottle_slug="demo", decision="approve")
|
|
self.assertTrue(ok, err)
|
|
# policy live-applied so /resolve serves the new routes
|
|
rec = self.store.get(bottle_id)
|
|
assert rec is not None
|
|
self.assertEqual(new_routes, rec.policy)
|
|
# response written -> agent unblocks, proposal no longer pending
|
|
self.assertEqual(STATUS_APPROVED, read_response("demo", pid).status)
|
|
self.assertEqual([], self.orch.supervise_pending())
|
|
|
|
def test_pending_carries_human_label(self) -> None:
|
|
# The proposal is keyed by bottle_id, but pending dicts also expose the
|
|
# bottle's human slug so the operator sees a name, not a hex id.
|
|
bottle_id = self._register("codex-dev-a1b2c", "routes: []\n")
|
|
self._queue(bottle_id, "routes:\n - host: google.com\n")
|
|
pending = self.orch.supervise_pending()
|
|
self.assertEqual(bottle_id, pending[0]["bottle_slug"])
|
|
self.assertEqual("codex-dev-a1b2c", pending[0]["bottle_label"])
|
|
|
|
def test_pending_label_falls_back_to_slug_when_bottle_gone(self) -> None:
|
|
# No registry record (torn down): label is the id, never empty.
|
|
self._queue("ghost-id", "routes:\n - host: google.com\n")
|
|
pending = self.orch.supervise_pending()
|
|
self.assertEqual("ghost-id", pending[0]["bottle_label"])
|
|
|
|
def test_approve_by_bottle_id_applies_policy(self) -> None:
|
|
# Consolidated reality: the supervise server keys each proposal by the
|
|
# orchestrator-assigned bottle_id, not the human slug. Approval must
|
|
# resolve the record by that id and apply the policy (regression for
|
|
# the "bottle <id> is no longer registered" 409).
|
|
bottle_id = self._register("codex-dev-a1b2c", "routes: []\n")
|
|
new_routes = "routes:\n - host: google.com\n"
|
|
pid = self._queue(bottle_id, new_routes)
|
|
ok, err = self.orch.supervise_respond(
|
|
pid, bottle_slug=bottle_id, decision="approve")
|
|
self.assertTrue(ok, err)
|
|
rec = self.store.get(bottle_id)
|
|
assert rec is not None
|
|
self.assertEqual(new_routes, rec.policy)
|
|
|
|
def test_modify_applies_final_file_not_proposed(self) -> None:
|
|
bottle_id = self._register("demo", "routes: []\n")
|
|
pid = self._queue("demo", "routes:\n - host: google.com\n")
|
|
edited = "routes:\n - host: example.com\n"
|
|
ok, _ = self.orch.supervise_respond(
|
|
pid, bottle_slug="demo", decision="modify", final_file=edited)
|
|
self.assertTrue(ok)
|
|
rec = self.store.get(bottle_id)
|
|
assert rec is not None
|
|
self.assertEqual(edited, rec.policy)
|
|
|
|
def test_reject_leaves_policy_unchanged(self) -> None:
|
|
bottle_id = self._register("demo", "routes:\n - host: existing.com\n")
|
|
pid = self._queue("demo", "routes:\n - host: google.com\n")
|
|
ok, _ = self.orch.supervise_respond(
|
|
pid, bottle_slug="demo", decision="reject", notes="no")
|
|
self.assertTrue(ok)
|
|
rec = self.store.get(bottle_id)
|
|
assert rec is not None
|
|
self.assertEqual("routes:\n - host: existing.com\n", rec.policy)
|
|
self.assertEqual("rejected", read_response("demo", pid).status)
|
|
|
|
def test_unknown_proposal_is_error(self) -> None:
|
|
ok, err = self.orch.supervise_respond(
|
|
"ghost", bottle_slug="demo", decision="approve")
|
|
self.assertFalse(ok)
|
|
self.assertIn("no such proposal", err)
|
|
|
|
def test_unknown_decision_is_error(self) -> None:
|
|
self._register("demo", "routes: []\n")
|
|
pid = self._queue("demo", "routes:\n - host: google.com\n")
|
|
ok, err = self.orch.supervise_respond(
|
|
pid, bottle_slug="demo", decision="bogus")
|
|
self.assertFalse(ok)
|
|
self.assertIn("unknown decision", err)
|
|
|
|
def test_approve_when_bottle_gone_cannot_apply(self) -> None:
|
|
# Proposal queued but the bottle was torn down before the operator
|
|
# acted: an egress apply has no target, so respond fails closed.
|
|
pid = self._queue("ghost-bottle", "routes:\n - host: google.com\n")
|
|
ok, err = self.orch.supervise_respond(
|
|
pid, bottle_slug="ghost-bottle", decision="approve")
|
|
self.assertFalse(ok)
|
|
self.assertIn("no longer registered", err)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|