Spaces:
Running
Running
| from __future__ import annotations | |
| import json | |
| import os | |
| import tempfile | |
| import unittest | |
| from concurrent.futures import ThreadPoolExecutor | |
| from dataclasses import replace | |
| from pathlib import Path | |
| from unittest.mock import patch | |
| from lumi_eggcracker.brokered import BrokeredOperator, BrokeredStoreError | |
| from lumi_eggcracker.brokered import operator as implementation | |
| def canonical(value: object) -> bytes: | |
| return (json.dumps(value, sort_keys=True, separators=(",", ":")) + "\n").encode() | |
| class BrokeredOperatorTests(unittest.TestCase): | |
| def setUp(self) -> None: | |
| self.temp = tempfile.TemporaryDirectory(dir=tempfile.gettempdir()) | |
| self.addCleanup(self.temp.cleanup) | |
| self.state_dir = Path(self.temp.name) / "state" | |
| self.registrar, self.operator = BrokeredOperator.bootstrap(self.state_dir) | |
| def grant_and_client(self): | |
| grant = self.registrar.register_run() | |
| return grant, self.operator.client(grant) | |
| def test_allowed_effect_replay_and_state_reopen(self) -> None: | |
| grant, client = self.grant_and_client() | |
| queued = client.admit("operation-1") | |
| self.assertEqual(("admission", "QUEUED", "ADMITTED"), (queued.phase, queued.outcome, queued.code)) | |
| self.assertNotIn("capability", queued.as_dict()) | |
| self.assertLessEqual(len(queued.canonical_bytes()), implementation.MAX_RECEIPT_BYTES) | |
| applied = self.operator.dispatch(queued.queue_id) | |
| self.assertTrue(applied.effect_applied) | |
| self.assertEqual("APPLIED", applied.outcome) | |
| self.assertEqual("REPLAY", self.operator.dispatch(queued.queue_id).code) | |
| with self.assertRaises(TypeError): | |
| self.operator.dispatch(queued.queue_id, target="synthetic.other") | |
| self.assertEqual( | |
| implementation.WorldSnapshot(1, implementation.UNRELATED_CANARY_ALLOCATION), | |
| self.operator.world_snapshot(), | |
| ) | |
| _, reopened = BrokeredOperator.open(self.state_dir) | |
| self.assertEqual(self.operator.world_snapshot(), reopened.world_snapshot()) | |
| self.assertEqual("REPLAY", reopened.dispatch(queued.queue_id).code) | |
| self.assertEqual("operation-1", json.loads(client.request_bytes("operation-1"))["operation_id"]) | |
| self.assertEqual(grant.run_id, applied.run_id) | |
| def test_parser_rejects_duplicate_unknown_boolean_malformed_oversized_and_noncanonical(self) -> None: | |
| _, client = self.grant_and_client() | |
| valid = client.request_bytes("strict-input") | |
| duplicate = valid.replace(b'"action":"increment",', b'"action":"increment","action":"increment",') | |
| unknown = json.loads(valid) | |
| unknown["label"] = "caller-controlled" | |
| boolean_generation = valid.replace(b'"generation":0', b'"generation":true') | |
| cases = ( | |
| (duplicate, "MALFORMED_INPUT"), | |
| (canonical(unknown), "UNKNOWN_FIELDS"), | |
| (boolean_generation, "MALFORMED_INPUT"), | |
| (b'{"broken":', "MALFORMED_INPUT"), | |
| (b" " * (implementation.MAX_REQUEST_BYTES + 1), "MALFORMED_INPUT"), | |
| (valid.rstrip(b"\n"), "NON_CANONICAL_INPUT"), | |
| ) | |
| for raw, code in cases: | |
| with self.subTest(code=code, size=len(raw)): | |
| result = self.operator.admit(raw) | |
| self.assertEqual(code, result.code) | |
| self.assertFalse(result.effect_applied) | |
| self.assertEqual(0, self.operator.world_snapshot().protected_effects) | |
| def test_target_run_and_generation_relabelling_do_not_authorize(self) -> None: | |
| grant, client = self.grant_and_client() | |
| other = self.registrar.register_run() | |
| with self.assertRaises(PermissionError): | |
| implementation.TrustedRegistrar(self.operator) | |
| with self.assertRaises(TypeError): | |
| self.registrar.register_run(run_id="caller-selected") | |
| with self.assertRaisesRegex(ValueError, "labels"): | |
| self.operator.client(replace(grant, target="synthetic.other")) | |
| valid = json.loads(client.request_bytes("relabel-1")) | |
| valid["run_id"] = other.run_id | |
| self.assertEqual("CAPABILITY_INVALID", self.operator.admit(canonical(valid)).code) | |
| valid = json.loads(client.request_bytes("relabel-2")) | |
| valid["generation"] = 77 | |
| self.assertEqual("CAPABILITY_INVALID", self.operator.admit(canonical(valid)).code) | |
| valid = json.loads(client.request_bytes("relabel-3")) | |
| valid["target"] = "synthetic.other" | |
| self.assertEqual("MALFORMED_INPUT", self.operator.admit(canonical(valid)).code) | |
| self.assertEqual(0, self.operator.world_snapshot().protected_effects) | |
| def test_expiry_is_exact_at_the_capability_boundary(self) -> None: | |
| with patch.object(implementation.time, "time_ns", return_value=1_000_000_000): | |
| grant = self.registrar.register_run() | |
| client = self.operator.client(grant) | |
| with patch.object( | |
| implementation.time, | |
| "time_ns", | |
| return_value=grant.expires_at_ms * 1_000_000, | |
| ): | |
| self.assertEqual("EXPIRED", client.admit("at-expiry").code) | |
| self.assertEqual(0, self.operator.world_snapshot().protected_effects) | |
| def test_per_run_and_aggregate_pending_budgets(self) -> None: | |
| _, first_client = self.grant_and_client() | |
| _, second_client = self.grant_and_client() | |
| first_request = first_client.request_bytes("budget-first") | |
| second_request = second_client.request_bytes("budget-other") | |
| self.assertEqual(len(first_request), len(second_request)) | |
| with patch.object( | |
| implementation, | |
| "MAX_PENDING_BYTES", | |
| (len(first_request) * 2) - 1, | |
| ): | |
| first = self.operator.admit(first_request) | |
| self.assertEqual("QUEUED", first.outcome) | |
| self.assertEqual("GLOBAL_BUDGET_EXHAUSTED", self.operator.admit(second_request).code) | |
| self.assertTrue(self.operator.dispatch(first.queue_id).effect_applied) | |
| _, client = self.grant_and_client() | |
| for index in range(implementation.MAX_ACTIONS_PER_RUN): | |
| queued = client.admit(f"run-budget-{index}") | |
| self.assertEqual("QUEUED", queued.outcome) | |
| self.operator.dispatch(queued.queue_id) | |
| self.assertEqual("BUDGET_EXHAUSTED", client.admit("run-budget-overflow").code) | |
| self.assertEqual( | |
| implementation.MAX_ACTIONS_PER_RUN + 1, | |
| self.operator.world_snapshot().protected_effects, | |
| ) | |
| def test_stop_request_fences_queued_and_future_actions_without_claiming_process_stop(self) -> None: | |
| grant, client = self.grant_and_client() | |
| queued = client.admit("before-stop") | |
| stop = self.operator.request_stop(grant.run_id) | |
| self.assertEqual(("revocation", "REVOKED", 1), (stop.phase, stop.outcome, stop.generation)) | |
| self.assertEqual("UNSUPPORTED", stop.process_stop) | |
| self.assertEqual("STALE_GENERATION", self.operator.dispatch(queued.queue_id).code) | |
| self.assertEqual("CAPABILITY_REVOKED", client.admit("after-stop").code) | |
| self.assertEqual("UNSUPPORTED", self.operator.verified_process_stop(grant.run_id).outcome) | |
| self.assertEqual(0, self.operator.world_snapshot().protected_effects) | |
| self.assertEqual(implementation.UNRELATED_CANARY_ALLOCATION, self.operator.world_snapshot().unrelated_canary_allocation) | |
| _, reopened = BrokeredOperator.open(self.state_dir) | |
| self.assertEqual(1, reopened.revoke(grant.run_id).generation) | |
| self.assertEqual("ALREADY_REVOKED", reopened.request_stop(grant.run_id).outcome) | |
| self.assertEqual("ALREADY_FINAL", reopened.dispatch(queued.queue_id).code) | |
| def test_concurrent_dispatch_has_one_synthetic_effect(self) -> None: | |
| _, client = self.grant_and_client() | |
| queued = client.admit("concurrent-1") | |
| with ThreadPoolExecutor(max_workers=2) as pool: | |
| results = list(pool.map(self.operator.dispatch, [queued.queue_id, queued.queue_id])) | |
| self.assertCountEqual(["EFFECT_APPLIED", "REPLAY"], [result.code for result in results]) | |
| self.assertEqual(1, self.operator.world_snapshot().protected_effects) | |
| self.assertEqual(implementation.UNRELATED_CANARY_ALLOCATION, self.operator.world_snapshot().unrelated_canary_allocation) | |
| def test_full_replay_registry_does_not_evict_live_generation_fence(self) -> None: | |
| with patch.object(implementation, "MAX_REPLAY_KEYS", 1): | |
| first_grant, first_client = self.grant_and_client() | |
| first = first_client.admit("one-only-key") | |
| _, second_client = self.grant_and_client() | |
| self.assertEqual("REPLAY_CAPACITY_EXHAUSTED", second_client.admit("second-key").code) | |
| self.assertEqual(1, self.operator.request_stop(first_grant.run_id).generation) | |
| _, reopened = BrokeredOperator.open(self.state_dir) | |
| self.assertEqual("STALE_GENERATION", reopened.dispatch(first.queue_id).code) | |
| self.assertEqual("REPLAY_CAPACITY_EXHAUSTED", reopened.client(self.registrar.register_run()).admit("third-key").code) | |
| self.assertEqual(0, reopened.world_snapshot().protected_effects) | |
| def test_state_missing_malformed_or_replaced_fails_closed(self) -> None: | |
| journal = self.state_dir / implementation.JOURNAL_NAME | |
| with self.subTest(case="missing"): | |
| journal.unlink() | |
| with self.assertRaises(BrokeredStoreError): | |
| BrokeredOperator.open(self.state_dir) | |
| self.registrar, self.operator = BrokeredOperator.bootstrap(self.state_dir.parent / "malformed") | |
| malformed = self.state_dir.parent / "malformed" | |
| malformed_journal = malformed / implementation.JOURNAL_NAME | |
| content = malformed_journal.read_bytes() | |
| malformed_journal.write_bytes(b"X" + content[1:]) | |
| with self.subTest(case="malformed"), self.assertRaises(BrokeredStoreError): | |
| BrokeredOperator.open(malformed) | |
| self.registrar, self.operator = BrokeredOperator.bootstrap(self.state_dir.parent / "replaced") | |
| replaced = self.state_dir.parent / "replaced" | |
| replaced_journal = replaced / implementation.JOURNAL_NAME | |
| replacement = replaced / "replacement.bin" | |
| replacement.write_bytes(replaced_journal.read_bytes()) | |
| os.replace(replacement, replaced_journal) | |
| with self.subTest(case="replaced"), self.assertRaises(BrokeredStoreError): | |
| BrokeredOperator.open(replaced) | |
| self.registrar, self.operator = BrokeredOperator.bootstrap(self.state_dir.parent / "rollback") | |
| rollback = self.state_dir.parent / "rollback" | |
| rollback_journal = rollback / implementation.JOURNAL_NAME | |
| prior_journal = rollback_journal.read_bytes() | |
| self.registrar.register_run() | |
| rollback_journal.write_bytes(prior_journal) | |
| with self.subTest(case="valid-prefix-rollback"), self.assertRaises(BrokeredStoreError): | |
| BrokeredOperator.open(rollback) | |
| if __name__ == "__main__": | |
| unittest.main() | |