"""Helper utilities for sandbox checkpoint Robot integration tests. Covers the M4 checkpoint infrastructure: - CheckpointManager create/rollback/list/delete - PlanExecutor checkpoint integration """ from __future__ import annotations import os import sys import tempfile from unittest.mock import MagicMock from cleveragents.infrastructure.sandbox.checkpoint import ( CheckpointManager, SandboxCheckpoint, ) from cleveragents.infrastructure.sandbox.no_sandbox import NoSandbox def _checkpoint_create() -> None: """Integration test: create a checkpoint and verify its fields.""" mgr = CheckpointManager() with tempfile.TemporaryDirectory() as tmpdir: # Create a file in the directory with open(os.path.join(tmpdir, "file.txt"), "w") as f: f.write("original content\n") sandbox = NoSandbox(resource_id="res-robot-01", original_path=tmpdir) sandbox.create(plan_id="plan-robot-01") cp = mgr.create_checkpoint( sandbox=sandbox, plan_id="plan-robot-01", phase="pre_execute", metadata={"sandbox_path": tmpdir}, ) assert isinstance(cp, SandboxCheckpoint) assert cp.phase == "pre_execute" assert cp.plan_id == "plan-robot-01" assert cp.sandbox_id == sandbox.sandbox_id assert len(cp.checkpoint_id) == 26 # ULID length assert cp.created_at is not None assert os.path.isdir(cp.snapshot_path) print("checkpoint-create-ok") def _checkpoint_rollback() -> None: """Integration test: create checkpoint, modify, rollback, verify.""" mgr = CheckpointManager() with tempfile.TemporaryDirectory() as tmpdir: # Create original file file_path = os.path.join(tmpdir, "data.txt") with open(file_path, "w") as f: f.write("original\n") sandbox = NoSandbox(resource_id="res-robot-02", original_path=tmpdir) sandbox.create(plan_id="plan-robot-02") # Checkpoint cp = mgr.create_checkpoint( sandbox=sandbox, plan_id="plan-robot-02", phase="pre_execute", metadata={"sandbox_path": tmpdir}, ) # Modify the file with open(file_path, "w") as f: f.write("modified\n") with open(file_path) as f: assert f.read() == "modified\n" # Add a new file new_file = os.path.join(tmpdir, "extra.txt") with open(new_file, "w") as f: f.write("extra\n") # Rollback result = mgr.rollback_to(cp) assert result is True # Verify original state restored with open(file_path) as f: content = f.read() assert content == "original\n", f"Expected 'original\\n', got {content!r}" # New file should be gone assert not os.path.exists(new_file), "extra.txt should not exist" print("checkpoint-rollback-ok") def _checkpoint_listing() -> None: """Integration test: multiple checkpoints maintain creation order.""" mgr = CheckpointManager() with tempfile.TemporaryDirectory() as tmpdir: sandbox = NoSandbox(resource_id="res-robot-03", original_path=tmpdir) sandbox.create(plan_id="plan-robot-03") phases = ["pre_execute", "post_execute", "pre_apply"] for phase in phases: mgr.create_checkpoint( sandbox=sandbox, plan_id="plan-robot-03", phase=phase, ) checkpoints = mgr.list_checkpoints(sandbox.sandbox_id) assert len(checkpoints) == 3 for i, phase in enumerate(phases): assert checkpoints[i].phase == phase # Verify chronological order for i in range(1, len(checkpoints)): assert checkpoints[i].created_at >= checkpoints[i - 1].created_at print("checkpoint-listing-ok") def _checkpoint_delete() -> None: """Integration test: delete a checkpoint and verify removal.""" mgr = CheckpointManager() with tempfile.TemporaryDirectory() as tmpdir: sandbox = NoSandbox(resource_id="res-robot-04", original_path=tmpdir) sandbox.create(plan_id="plan-robot-04") cp = mgr.create_checkpoint( sandbox=sandbox, plan_id="plan-robot-04", phase="pre_execute", ) assert len(mgr.list_checkpoints(sandbox.sandbox_id)) == 1 result = mgr.delete_checkpoint(cp.checkpoint_id) assert result is True assert len(mgr.list_checkpoints(sandbox.sandbox_id)) == 0 # Deleting again should return False result2 = mgr.delete_checkpoint(cp.checkpoint_id) assert result2 is False print("checkpoint-delete-ok") def _executor_integration() -> None: """Integration test: PlanExecutor accepts optional checkpoint_manager.""" from cleveragents.application.services.plan_executor import PlanExecutor lifecycle = MagicMock() # Without checkpoint manager executor_no_cp = PlanExecutor(lifecycle_service=lifecycle) assert executor_no_cp.checkpoint_manager is None # With checkpoint manager mgr = CheckpointManager() executor_with_cp = PlanExecutor( lifecycle_service=lifecycle, checkpoint_manager=mgr, ) assert executor_with_cp.checkpoint_manager is mgr print("executor-integration-ok") def main() -> None: if len(sys.argv) < 2: raise SystemExit("Expected command argument") command = sys.argv[1] commands = { "checkpoint-create": _checkpoint_create, "checkpoint-rollback": _checkpoint_rollback, "checkpoint-listing": _checkpoint_listing, "checkpoint-delete": _checkpoint_delete, "executor-integration": _executor_integration, } if command not in commands: raise SystemExit(f"Unknown command: {command}") commands[command]() if __name__ == "__main__": main()