cleveragents-core/robot/helper_sandbox_checkpoint.py

"""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()