import base64 import importlib.util import io import json from pathlib import Path import tempfile import threading import unittest from unittest.mock import MagicMock, call, mock_open, patch spec = importlib.util.spec_from_file_location("ec_agent", Path(__file__).resolve().parents[1] / "ec_agent.py") module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) Agent, AgentError, Held, Journal = module.Agent, module.AgentError, module.Held, module.Journal INSTANCE = "a53d99da-a1eb-4f01-9b10-a6268b556542" class FakeDocker: def __init__(self): self.containers = {} self.calls = [] self.fail_restart = False self.lose_reply_after = None self.started = 0 def inspect(self, name): return self.containers.get(name) def call(self, *args, **kwargs): self.calls.append(args) if args[0] == "info": return json.dumps({"ServerVersion": "27.5.1"}) if args[:2] == ("image", "ls"): return "ghcr.io/ayymoss/plutainer:latest\n" if args[:2] == ("container", "create"): labels = {args[index + 1].split("=", 1)[0]: args[index + 1].split("=", 1)[1] for index, arg in enumerate(args) if arg == "--label"} name = args[args.index("--name") + 1] self.containers[name] = {"Id": "container-123", "Config": {"Labels": labels}, "State": {"Status": "created", "Running": False}} if args[:2] in (("container", "start"), ("container", "restart")): if self.fail_restart and args[1] == "restart": raise AgentError("Docker command unavailable or timed out") self.started += 1 self.containers[args[-1]]["State"] = {"Status": "running", "Running": True, "StartedAt": f"started-{self.started}"} if args[:2] == ("container", "stop"): self.containers[args[-1]]["State"] = {"Status": "exited", "Running": False} if args[:2] == ("container", "rm"): self.containers.pop(args[-1]) if args[0] == "container" and args[1] == self.lose_reply_after: raise AgentError("Docker command unavailable or timed out") return "" class Response(io.BytesIO): def __init__(self, content, status=200, headers=None): super().__init__(content) self.status = status self.headers = headers or {} class JournalRetentionTests(unittest.TestCase): def setUp(self): self.temp = tempfile.TemporaryDirectory() self.path = Path(self.temp.name) / "operations.sqlite3" self.journal = Journal(self.path) self.now = 2000000 self.sequence = 0 def tearDown(self): self.journal.close() self.temp.cleanup() def record(self, kind="files.stat", payload=None, result=None, state="completed", error=None, at=None, confirmed=True): self.sequence += 1 operation = {"id": f"operation-{self.sequence}", "kind": kind, "payload": payload or {}} self.journal.accept(operation) with patch.object(module.time, "time", return_value=self.now - 172800 if at is None else at): row = self.journal.update(operation["id"], state, result, error) if confirmed: self.journal.confirm_completed(row) return row def prune(self, days=1): with patch.object(module.time, "time", return_value=self.now): return self.journal.prune_observations(days) def test_cutoff_requires_an_older_confirmed_success(self): old = self.record(at=self.now - 86401) boundary = self.record(at=self.now - 86400) recent = self.record(at=self.now - 10) unconfirmed = self.record(confirmed=False) self.assertEqual(self.prune(), 1) with self.assertRaisesRegex(AgentError, "Unknown operation"): self.journal.get(old["id"]) for row in (boundary, recent, unconfirmed): self.assertEqual(self.journal.get(row["id"])["state"], "completed") def test_zero_and_invalid_retention_preserve_every_record(self): row = self.record() for days in (0, -1, 36501, True, False, None, "1", 1.0, 1.5): with self.subTest(days=days): self.assertEqual(self.prune(days), 0) self.assertEqual(self.journal.get(row["id"])["state"], "completed") def test_all_read_kinds_are_eligible_but_mutations_keep_their_replay_journal(self): for kind in ("host.inspect", "instance.inspect", "files.read", "files.list", "files.stat", "files.manifest", "logs.read"): self.record(kind, result={"healthy": True}) for command in ("status", "getstatus", "version"): self.record("rcon.execute", {"command": command}, {"ok": True}) mutations = [self.record(kind) for kind in ("instance.restart", "instance.stop", "instance.create", "files.stage", "files.activate", "credentials.rotate")] mutations.append(self.record("rcon.execute", {"command": "map_restart"}, {"ok": True})) self.assertEqual(self.prune(), 10) for row in mutations: saved, accepted = self.journal.accept(json.loads(row["request"])) self.assertFalse(accepted) self.assertEqual(saved["id"], row["id"]) def test_errors_failures_and_unfinished_operations_are_retained(self): retained = [self.record(result=result) for result in ({"ok": False}, {"healthy": False}, {"error": "partial result"}, {"error": False}, {"error": []})] retained.append(self.record(error="read did not finish")) retained.extend(self.record(state=state) for state in ("accepted", "running", "unknown", "failed")) self.assertEqual(self.prune(), 0) for row in retained: self.assertEqual(self.journal.get(row["id"])["state"], row["state"]) def test_old_receipt_cannot_confirm_a_changed_or_reexecuted_result(self): old = self.record(confirmed=False) with patch.object(module.time, "time", return_value=self.now - 90000): current = self.journal.update(old["id"], "completed", {"new": "result"}) self.journal.confirm_completed(old) self.assertIsNone(self.journal.get(old["id"])["reported_at"]) self.assertEqual(self.prune(), 0) self.journal.confirm_completed(current) self.assertEqual(self.prune(), 1) def test_pruning_uses_bounded_indexed_batches(self): template = self.record() self.journal.db.executemany("""INSERT INTO operations(id,fingerprint,request,state,result,updated,reported_at,retention_eligible) VALUES(?,?,?,'completed',NULL,?,?,1)""", [ (f"batch-{index}", template["fingerprint"], template["request"], self.now - 172800, self.now - 172800) for index in range(504) ]) self.journal.db.commit() self.assertEqual(self.prune(), 500) self.assertEqual(self.journal.db.execute("SELECT COUNT(*) FROM operations").fetchone()[0], 5) self.assertEqual(self.prune(), 5) plan = self.journal.db.execute("""EXPLAIN QUERY PLAN SELECT id FROM operations WHERE state='completed' AND reported_at IS NOT NULL AND retention_eligible=1 AND updated