bot-bottle/tests/unit/test_orchestrator_service.py

"""Unit tests for the Orchestrator launch lifecycle (PRD 0070)."""

from __future__ import annotations

import json
import secrets
import sqlite3
import tempfile
import unittest
from contextlib import closing
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.orchestrator.secret_store import new_env_var_secret
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_encrypted_tokens_can_be_reprovisioned_after_memory_loss(self) -> None:
        key = new_env_var_secret()
        rec = self.orch.launch_bottle(
            "10.243.0.12", tokens={"EGRESS_TOKEN_0": "secret"},
            env_var_secret=key,
        )
        self.assertNotEqual({}, self.store.get_agent_secrets(rec.bottle_id))
        self.orch._tokens.clear()
        self.assertTrue(self.orch.reprovision_from_secret(rec.bottle_id, key))
        self.assertEqual({"EGRESS_TOKEN_0": "secret"}, self.orch.tokens_for(rec.bottle_id))

    def test_reprovision_rejects_missing_rows_and_wrong_key(self) -> None:
        self.assertFalse(self.orch.reprovision_from_secret("missing", new_env_var_secret()))
        rec = self.orch.launch_bottle(
            "10.243.0.13", tokens={"K": "value"},
            env_var_secret=new_env_var_secret(),
        )
        self.assertFalse(self.orch.reprovision_from_secret(rec.bottle_id, new_env_var_secret()))

    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)

    # --- agent half: queue + poll (issue #469) -----------------------------

    def test_queue_proposal_then_poll_pending(self) -> None:
        bottle_id = self._register("demo", "routes: []\n")
        pid = self.orch.supervise_queue_proposal(
            bottle_id, tool=TOOL_EGRESS_ALLOW,
            proposed_file="routes:\n  - host: google.com\n", justification="need it")
        # Visible to the operator, keyed by the bottle id.
        pending = self.orch.supervise_pending()
        self.assertEqual([pid], [p["id"] for p in pending])
        self.assertEqual(
            {"status": "pending"}, self.orch.supervise_poll_response(bottle_id, pid))

    def test_poll_returns_decision_and_archives(self) -> None:
        bottle_id = self._register("demo", "routes: []\n")
        pid = self.orch.supervise_queue_proposal(
            bottle_id, tool=TOOL_EGRESS_ALLOW,
            proposed_file="routes:\n  - host: google.com\n", justification="need it")
        self.orch.supervise_respond(
            pid, bottle_slug=bottle_id, decision="approve", notes="ok")
        decided = self.orch.supervise_poll_response(bottle_id, pid)
        self.assertEqual("approved", decided["status"])
        self.assertEqual("ok", decided["notes"])
        # Archived on read → gone from pending, and a re-poll is 'unknown'.
        self.assertEqual([], self.orch.supervise_pending())
        self.assertEqual(
            {"status": "unknown"}, self.orch.supervise_poll_response(bottle_id, pid))

    def test_poll_unknown_for_other_bottle(self) -> None:
        bottle_id = self._register("demo", "routes: []\n")
        pid = self.orch.supervise_queue_proposal(
            bottle_id, tool=TOOL_EGRESS_ALLOW,
            proposed_file="routes:\n  - host: google.com\n", justification="j")
        # A different bottle id can't read demo's proposal (scoped by queue key).
        self.assertEqual(
            {"status": "unknown"}, self.orch.supervise_poll_response("other-bottle", pid))


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


class TestOrchestratorReconcile(unittest.TestCase):
    """`reconcile` — drop rows for bottles that are no longer running."""

    def setUp(self) -> None:
        self._tmp = tempfile.TemporaryDirectory()
        self.secret = secrets.token_bytes(16)
        self.db = Path(self._tmp.name) / "r.db"
        self.store = RegistryStore(self.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 _age_all(self, seconds: float) -> None:
        """Backdate every row past the reap grace window."""
        with closing(sqlite3.connect(self.db)) as conn:
            conn.execute(
                "UPDATE orchestrator_bottles SET created_at = created_at - ?", (seconds,))
            conn.commit()

    def test_reaps_dead_bottle_and_forgets_its_tokens(self) -> None:
        dead = self.orch.launch_bottle("10.243.0.1", tokens={"EGRESS_TOKEN_0": "s3cret"})
        live = self.orch.launch_bottle("10.243.0.2", tokens={"EGRESS_TOKEN_0": "keep"})
        self._age_all(600)

        self.assertEqual([dead.bottle_id], self.orch.reconcile(["10.243.0.2"]))
        self.assertIsNone(self.store.get(dead.bottle_id))
        self.assertIsNotNone(self.store.get(live.bottle_id))
        # The in-memory egress credential goes with the row.
        self.assertEqual({}, self.orch.tokens_for(dead.bottle_id))
        self.assertEqual({"EGRESS_TOKEN_0": "keep"}, self.orch.tokens_for(live.bottle_id))

    def test_reconcile_does_not_broker_a_teardown(self) -> None:
        """The container is already gone — there is nothing to stop, and a
        broker error must not stop the sweep clearing the row."""
        self.orch.launch_bottle("10.243.0.1")
        self._age_all(600)
        self.broker.launched.clear()
        self.orch.reconcile([])
        self.assertEqual([], self.broker.torn_down)

    def test_reconcile_keeps_everything_when_all_are_live(self) -> None:
        a = self.orch.launch_bottle("10.243.0.1")
        b = self.orch.launch_bottle("10.243.0.2")
        self._age_all(600)
        self.assertEqual([], self.orch.reconcile(["10.243.0.1", "10.243.0.2"]))
        self.assertIsNotNone(self.store.get(a.bottle_id))
        self.assertIsNotNone(self.store.get(b.bottle_id))