"""Unit tests for the orchestrator bottle registry + attribution (PRD 0070).""" from __future__ import annotations 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 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. a = self.store.register("10.243.0.1", bottle_id="a") self.store.register("10.243.0.1", bottle_id="b") self.assertIsNone(self.store.attribute("10.243.0.1", a.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.store.register("10.243.0.1", bottle_id="b") self.assertIsNone(self.store.by_source_ip("10.243.0.1")) # ambiguous if __name__ == "__main__": unittest.main()