bot-bottle/tests/unit/test_orchestrator_registry.py

"""Unit tests for the orchestrator bottle registry + attribution (PRD 0070)."""

from __future__ import annotations

import sqlite3
import tempfile
import unittest
from pathlib import Path

from bot_bottle.orchestrator.registry import (
    BottleRecord,
    RegistryStore,
    new_identity_token,
)


class TestRegistryStore(unittest.TestCase):
    def setUp(self) -> None:
        self._tmp = tempfile.TemporaryDirectory()
        self.db = Path(self._tmp.name) / "registry.db"
        self.store = RegistryStore(self.db)
        self.store.migrate()

    def tearDown(self) -> None:
        self._tmp.cleanup()

    def _force_active_row(self, source_ip: str, bottle_id: str) -> None:
        """Insert a raw active row, bypassing `register`'s same-IP supersede —
        the only way to manufacture the ambiguity the fail-closed guard exists
        for (crash-orphaned / externally-injected duplicate)."""
        conn = sqlite3.connect(self.db)
        conn.execute(
            "INSERT INTO orchestrator_bottles "
            "(bottle_id, source_ip, identity_token, state, created_at, metadata, policy) "
            "VALUES (?, ?, ?, 'active', 0.0, '', '')",
            (bottle_id, source_ip, "tok-" + bottle_id),
        )
        conn.commit()
        conn.close()

    def test_register_mints_id_and_token(self) -> None:
        rec = self.store.register("10.243.0.1")
        self.assertTrue(rec.bottle_id)
        self.assertTrue(rec.identity_token)
        self.assertEqual("active", rec.state)
        self.assertEqual(rec, self.store.get(rec.bottle_id))

    def test_identity_tokens_are_unique(self) -> None:
        tokens = {new_identity_token() for _ in range(200)}
        self.assertEqual(200, len(tokens))
        a = self.store.register("10.243.0.1")
        b = self.store.register("10.243.0.3")
        self.assertNotEqual(a.identity_token, b.identity_token)

    def test_all_and_deregister(self) -> None:
        a = self.store.register("10.243.0.1")
        b = self.store.register("10.243.0.3")
        self.assertEqual({a.bottle_id, b.bottle_id}, {r.bottle_id for r in self.store.all()})
        self.assertTrue(self.store.deregister(a.bottle_id))
        self.assertEqual([b.bottle_id], [r.bottle_id for r in self.store.all()])

    def test_deregister_missing_is_false(self) -> None:
        self.assertFalse(self.store.deregister("nope"))

    def test_explicit_id_replaces(self) -> None:
        self.store.register("10.243.0.1", bottle_id="fixed", metadata="a")
        self.store.register("10.243.0.9", bottle_id="fixed", metadata="b")
        rec = self.store.get("fixed")
        assert rec is not None
        self.assertEqual("10.243.0.9", rec.source_ip)
        self.assertEqual("b", rec.metadata)
        self.assertEqual(1, len(self.store.all()))

    def test_attribute_success_requires_ip_and_token(self) -> None:
        rec = self.store.register("10.243.0.1")
        got = self.store.attribute("10.243.0.1", rec.identity_token)
        assert got is not None
        self.assertEqual(rec.bottle_id, got.bottle_id)

    def test_attribute_wrong_token_denied(self) -> None:
        self.store.register("10.243.0.1")
        self.assertIsNone(self.store.attribute("10.243.0.1", "wrong-token"))

    def test_attribute_unknown_ip_denied(self) -> None:
        rec = self.store.register("10.243.0.1")
        self.assertIsNone(self.store.attribute("10.243.9.9", rec.identity_token))

    def test_attribute_empty_token_denied(self) -> None:
        self.store.register("10.243.0.1")
        self.assertIsNone(self.store.attribute("10.243.0.1", ""))

    def test_attribute_ambiguous_ip_denied(self) -> None:
        # Two active bottles on one source IP is a misconfiguration — deny
        # rather than guess (fail-closed), even with a valid token. (Forced
        # directly: `register` now supersedes, so this can only arise from an
        # orphaned/injected row.)
        a = self.store.register("10.243.0.1", bottle_id="a")
        self._force_active_row("10.243.0.1", "b")
        self.assertIsNone(self.store.attribute("10.243.0.1", a.identity_token))

    def test_reregister_same_source_ip_supersedes(self) -> None:
        # Re-registering a source IP (reused IP / crash recovery / persisted
        # DB) must supersede the prior active bottle, not add a second active
        # row that would fail-closed the whole IP to 403.
        a = self.store.register("10.243.0.1", bottle_id="a", policy="A")
        b = self.store.register("10.243.0.1", bottle_id="b", policy="B")
        # exactly one active row at that IP, and it's the new bottle
        got = self.store.by_source_ip("10.243.0.1")
        assert got is not None
        self.assertEqual("b", got.bottle_id)
        self.assertEqual("B", got.policy)
        # the superseded bottle is gone and its token no longer attributes
        self.assertIsNone(self.store.get("a"))
        self.assertIsNone(self.store.attribute("10.243.0.1", a.identity_token))
        self.assertIsNotNone(self.store.attribute("10.243.0.1", b.identity_token))

    def test_state_persists_across_reopen(self) -> None:
        rec = self.store.register("10.243.0.1")
        reopened = RegistryStore(self.db)
        got = reopened.get(rec.bottle_id)
        assert got is not None
        self.assertEqual(rec, got)
        # Attribution works against the reopened (durable) store too.
        self.assertIsNotNone(reopened.attribute("10.243.0.1", rec.identity_token))

    def test_redacted_hides_token(self) -> None:
        rec = BottleRecord(
            bottle_id="x", source_ip="10.243.0.1", identity_token="secret"
        )
        self.assertNotIn("identity_token", rec.redacted())
        self.assertEqual("10.243.0.1", rec.redacted()["source_ip"])

    def test_register_with_policy_resolves_via_attribution(self) -> None:
        rec = self.store.register("10.243.0.1", policy='{"allow":["x"]}')
        self.assertEqual('{"allow":["x"]}', rec.policy)
        got = self.store.attribute("10.243.0.1", rec.identity_token)
        assert got is not None
        self.assertEqual('{"allow":["x"]}', got.policy)

    def test_set_policy_updates_live(self) -> None:
        rec = self.store.register("10.243.0.1")
        self.assertEqual("", rec.policy)
        self.assertTrue(self.store.set_policy(rec.bottle_id, '{"v":2}'))
        got = self.store.get(rec.bottle_id)
        assert got is not None
        self.assertEqual('{"v":2}', got.policy)

    def test_set_policy_missing_is_false(self) -> None:
        self.assertFalse(self.store.set_policy("ghost", "{}"))

    def test_policy_defaults_empty_and_persists(self) -> None:
        rec = self.store.register("10.243.0.1")
        got = RegistryStore(self.db).get(rec.bottle_id)
        assert got is not None
        self.assertEqual("", got.policy)

    def test_by_source_ip_returns_single_active(self) -> None:
        rec = self.store.register("10.243.0.1")
        got = self.store.by_source_ip("10.243.0.1")
        assert got is not None
        self.assertEqual(rec.bottle_id, got.bottle_id)

    def test_by_source_ip_unknown_or_ambiguous_denied(self) -> None:
        self.assertIsNone(self.store.by_source_ip("10.243.9.9"))  # unknown
        self.store.register("10.243.0.1", bottle_id="a")
        self._force_active_row("10.243.0.1", "b")  # orphaned duplicate
        self.assertIsNone(self.store.by_source_ip("10.243.0.1"))  # ambiguous


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