cleveragents-core/features/steps/checkpoint_manager_coverage_steps.py

"""Step definitions for checkpoint_manager_coverage.feature.

Targets every method, branch, and edge case in checkpoint.py to achieve
full line-rate and branch-rate coverage.
"""

from __future__ import annotations

import os
import shutil
import tempfile
import threading
from datetime import datetime
from typing import Any
from unittest.mock import MagicMock, patch

from behave import given, then, when
from behave.runner import Context

from cleveragents.infrastructure.sandbox.checkpoint import (
    Checkpointable,
    CheckpointManager,
    SandboxCheckpoint,
)
from cleveragents.infrastructure.sandbox.protocol import SandboxError

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_mock_sandbox(
    sandbox_id: str = "sb-mock-001",
    context: Any = None,
) -> MagicMock:
    """Build a mock that satisfies the Checkpointable protocol."""
    sb = MagicMock()
    sb.sandbox_id = sandbox_id
    sb.context = context
    return sb


def _make_temp_sandbox_dir() -> str:
    """Create a temp dir with sample files and a subdirectory."""
    tmpdir = tempfile.mkdtemp(prefix="ca-cpmgr-test-")
    os.makedirs(os.path.join(tmpdir, "subdir"), exist_ok=True)
    with open(os.path.join(tmpdir, "file.txt"), "w") as f:
        f.write("original\n")
    with open(os.path.join(tmpdir, "subdir", "nested.txt"), "w") as f:
        f.write("nested-original\n")
    return tmpdir


def _register_cleanup(context: Context, path: str) -> None:
    """Schedule a directory for cleanup after the scenario."""
    if not hasattr(context, "_cleanup_handlers"):
        context._cleanup_handlers = []
    context._cleanup_handlers.append(lambda: shutil.rmtree(path, ignore_errors=True))


# ===========================================================================
# SandboxCheckpoint Pydantic model
# ===========================================================================


@given("I build a SandboxCheckpoint with valid fields")
def step_build_checkpoint_model(context: Context) -> None:
    context.cp_model = SandboxCheckpoint(
        checkpoint_id="cp-aaa",
        sandbox_id="sb-bbb",
        plan_id="plan-ccc",
        phase="pre_execute",
        created_at=datetime(2025, 1, 1, 12, 0, 0),
        metadata={"key": "val"},
        snapshot_path="/tmp/snap",
    )


@then("the constructed checkpoint model should expose all field values")
def step_assert_checkpoint_fields(context: Context) -> None:
    cp = context.cp_model
    assert cp.checkpoint_id == "cp-aaa"
    assert cp.sandbox_id == "sb-bbb"
    assert cp.plan_id == "plan-ccc"
    assert cp.phase == "pre_execute"
    assert cp.created_at == datetime(2025, 1, 1, 12, 0, 0)
    assert cp.metadata == {"key": "val"}
    assert cp.snapshot_path == "/tmp/snap"


@then("the constructed checkpoint model should be frozen")
def step_assert_checkpoint_frozen(context: Context) -> None:
    try:
        context.cp_model.phase = "modified"
        raise AssertionError("Expected ValidationError for frozen model")
    except Exception:
        pass  # pydantic frozen model rejects mutation


@given("I build a SandboxCheckpoint without explicit metadata")
def step_build_checkpoint_no_metadata(context: Context) -> None:
    context.cp_model = SandboxCheckpoint(
        checkpoint_id="cp-ddd",
        sandbox_id="sb-eee",
        plan_id="plan-fff",
        phase="post_execute",
        created_at=datetime.now(),
        snapshot_path="/tmp/snap2",
    )


@then("the constructed checkpoint metadata should be an empty dict")
def step_assert_empty_metadata_default(context: Context) -> None:
    assert context.cp_model.metadata == {}


# ===========================================================================
# Checkpointable protocol
# ===========================================================================


@given("a mock object with sandbox_id and context properties")
def step_mock_checkpointable(context: Context) -> None:
    class _Good:
        @property
        def sandbox_id(self) -> str:
            return "sb-proto"

        @property
        def context(self) -> Any:
            return None

    context.proto_obj = _Good()


@then("the mock object should satisfy the Checkpointable protocol")
def step_assert_satisfies_checkpointable(context: Context) -> None:
    assert isinstance(context.proto_obj, Checkpointable)


@given("a mock object without a sandbox_id property")
def step_mock_not_checkpointable(context: Context) -> None:
    class _Bad:
        @property
        def context(self) -> Any:
            return None

    context.proto_obj = _Bad()


@then("the mock object should not satisfy the Checkpointable protocol")
def step_assert_not_checkpointable(context: Context) -> None:
    assert not isinstance(context.proto_obj, Checkpointable)


# ===========================================================================
# CheckpointManager.__init__
# ===========================================================================


@given("a freshly created CheckpointManager")
def step_fresh_manager(context: Context) -> None:
    context.mgr = CheckpointManager()
    context.crafted_cp = None
    context.cp_result = None
    context.rollback_result = None
    context.delete_result = None
    context.cp_list = None
    context.snapshot_result = None
    context.raised_error = None
    context.cleanup_error = None


@then("the manager internal checkpoints dict should be empty")
def step_assert_empty_checkpoints(context: Context) -> None:
    assert context.mgr._checkpoints == {}


@then("the manager should have a threading lock")
def step_assert_has_lock(context: Context) -> None:
    assert isinstance(context.mgr._lock, type(threading.RLock()))


# ===========================================================================
# create_checkpoint - context is None
# ===========================================================================


@given("a mock sandbox whose context is None")
def step_mock_sandbox_no_context(context: Context) -> None:
    context.mock_sb = _make_mock_sandbox(sandbox_id="sb-nocontext")


@when(
    'I invoke create_checkpoint with plan "{plan_id}" phase "{phase}" and no metadata'
)
def step_create_cp_no_metadata(context: Context, plan_id: str, phase: str) -> None:
    context.cp_result = context.mgr.create_checkpoint(
        sandbox=context.mock_sb,
        plan_id=plan_id,
        phase=phase,
    )


@then("the returned checkpoint should have a non-empty snapshot_path directory")
def step_assert_snapshot_path_dir(context: Context) -> None:
    sp = context.cp_result.snapshot_path
    assert sp, "snapshot_path should not be empty"
    assert os.path.isdir(sp), f"Expected directory at {sp}"
    _register_cleanup(context, os.path.dirname(sp))


@then("the returned checkpoint sandbox_id should equal the mock sandbox id")
def step_assert_cp_sandbox_id(context: Context) -> None:
    assert context.cp_result.sandbox_id == context.mock_sb.sandbox_id


@when(
    'I invoke create_checkpoint with plan "{plan_id}" phase "{phase}" and numeric metadata'
)
def step_create_cp_numeric_metadata(context: Context, plan_id: str, phase: str) -> None:
    context.cp_result = context.mgr.create_checkpoint(
        sandbox=context.mock_sb,
        plan_id=plan_id,
        phase=phase,
        metadata={"count": 42, "flag": True},
    )
    _register_cleanup(context, os.path.dirname(context.cp_result.snapshot_path))


@then("the returned checkpoint metadata values should all be strings")
def step_assert_metadata_strings(context: Context) -> None:
    for k, v in context.cp_result.metadata.items():
        assert isinstance(v, str), f"metadata[{k!r}] = {v!r} is not a string"


@when(
    'I invoke create_checkpoint with plan "{plan_id}" phase "{phase}" and None metadata'
)
def step_create_cp_none_metadata(context: Context, plan_id: str, phase: str) -> None:
    context.cp_result = context.mgr.create_checkpoint(
        sandbox=context.mock_sb,
        plan_id=plan_id,
        phase=phase,
        metadata=None,
    )
    _register_cleanup(context, os.path.dirname(context.cp_result.snapshot_path))


@then("the returned checkpoint metadata should be an empty dict")
def step_assert_cp_metadata_empty(context: Context) -> None:
    assert context.cp_result.metadata == {}


# ===========================================================================
# create_checkpoint - context with sandbox_path
# ===========================================================================


@given("a temporary sandbox directory with sample files")
def step_temp_sandbox_dir(context: Context) -> None:
    context.temp_sandbox = _make_temp_sandbox_dir()
    _register_cleanup(context, context.temp_sandbox)


@given("a mock sandbox whose context points to the temporary directory")
def step_mock_sandbox_with_context(context: Context) -> None:
    ctx = MagicMock()
    ctx.sandbox_path = context.temp_sandbox
    context.mock_sb = _make_mock_sandbox(sandbox_id="sb-withctx", context=ctx)


@then("the snapshot directory should contain copies of the sample files")
def step_assert_snapshot_has_files(context: Context) -> None:
    snap = context.cp_result.snapshot_path
    _register_cleanup(context, os.path.dirname(snap))
    assert os.path.isfile(os.path.join(snap, "file.txt"))
    assert os.path.isfile(os.path.join(snap, "subdir", "nested.txt"))
    with open(os.path.join(snap, "file.txt")) as f:
        assert f.read() == "original\n"


# ===========================================================================
# _snapshot_directory - branches
# ===========================================================================


@when("I call _snapshot_directory with a None sandbox_path")
def step_snapshot_none_path(context: Context) -> None:
    context.snapshot_result = context.mgr._snapshot_directory(None, "cp-test-none")
    _register_cleanup(context, os.path.dirname(context.snapshot_result))


@then("the returned snapshot should be an existing empty directory")
def step_assert_snapshot_empty_dir(context: Context) -> None:
    sp = context.snapshot_result
    assert os.path.isdir(sp), f"Expected directory at {sp}"
    assert os.listdir(sp) == [], f"Expected empty dir, got {os.listdir(sp)}"


@when("I call _snapshot_directory with a non-existent sandbox_path")
def step_snapshot_nonexistent_path(context: Context) -> None:
    context.snapshot_result = context.mgr._snapshot_directory(
        "/nonexistent/path/xyz", "cp-test-noexist"
    )
    _register_cleanup(context, os.path.dirname(context.snapshot_result))


@when("I call _snapshot_directory with the temporary sandbox_path")
def step_snapshot_valid_path(context: Context) -> None:
    context.snapshot_result = context.mgr._snapshot_directory(
        context.temp_sandbox, "cp-test-valid"
    )
    _register_cleanup(context, os.path.dirname(context.snapshot_result))


@then("the returned snapshot should contain the sample files")
def step_assert_snapshot_has_sample_files(context: Context) -> None:
    sp = context.snapshot_result
    assert os.path.isfile(os.path.join(sp, "file.txt"))
    assert os.path.isfile(os.path.join(sp, "subdir", "nested.txt"))


@when("I call _snapshot_directory with a path that triggers OSError")
def step_snapshot_oserror(context: Context) -> None:
    context.raised_error = None
    with patch("shutil.copytree", side_effect=OSError("mock copy failure")):
        try:
            context.mgr._snapshot_directory(context.temp_sandbox, "cp-test-err")
        except SandboxError as exc:
            context.raised_error = exc


@then("a SandboxError should have been raised")
def step_assert_sandbox_error(context: Context) -> None:
    assert context.raised_error is not None, "Expected SandboxError to be raised"
    assert isinstance(context.raised_error, SandboxError)
    assert "Failed to snapshot" in str(context.raised_error)


# ===========================================================================
# _cleanup_snapshot - branches
# ===========================================================================


@when("I call _cleanup_snapshot with an empty string")
def step_cleanup_empty_string(context: Context) -> None:
    context.cleanup_error = None
    try:
        CheckpointManager._cleanup_snapshot("")
    except Exception as exc:
        context.cleanup_error = exc


@then("no error should occur from cleanup_snapshot")
def step_assert_no_cleanup_error(context: Context) -> None:
    assert context.cleanup_error is None, (
        f"Expected no error, got {context.cleanup_error!r}"
    )


@given("a temporary snapshot directory tree for cleanup")
def step_temp_snapshot_for_cleanup(context: Context) -> None:
    parent = tempfile.mkdtemp(prefix="ca-cleanup-test-")
    snap = os.path.join(parent, "snapshot")
    os.makedirs(snap, exist_ok=True)
    with open(os.path.join(snap, "data.txt"), "w") as f:
        f.write("to-be-cleaned\n")
    context.cleanup_parent = parent
    context.cleanup_snapshot = snap


@when("I call _cleanup_snapshot with the temporary snapshot path")
def step_cleanup_real_snapshot(context: Context) -> None:
    context.cleanup_error = None
    try:
        CheckpointManager._cleanup_snapshot(context.cleanup_snapshot)
    except Exception as exc:
        context.cleanup_error = exc


@then("the temporary snapshot parent should no longer exist")
def step_assert_parent_removed(context: Context) -> None:
    assert not os.path.exists(context.cleanup_parent), (
        f"Parent should have been removed: {context.cleanup_parent}"
    )


@when("I call _cleanup_snapshot with a non-existent path")
def step_cleanup_nonexistent(context: Context) -> None:
    context.cleanup_error = None
    try:
        CheckpointManager._cleanup_snapshot("/nonexistent/snapshot/path")
    except Exception as exc:
        context.cleanup_error = exc


# ===========================================================================
# rollback_to - all branches
# ===========================================================================


@given("a SandboxCheckpoint with empty snapshot_path")
def step_cp_empty_snapshot(context: Context) -> None:
    context.crafted_cp = SandboxCheckpoint(
        checkpoint_id="cp-empty-snap",
        sandbox_id="sb-rb1",
        plan_id="plan-rb1",
        phase="pre_execute",
        created_at=datetime.now(),
        metadata={},
        snapshot_path="",
    )


@when("I invoke rollback_to on the crafted checkpoint")
def step_rollback_crafted(context: Context) -> None:
    if getattr(context, "rollback_should_patch_oserror", False):
        # Patch os.listdir to raise on the second call (copy phase)
        orig_listdir = os.listdir
        call_count = {"n": 0}

        def _failing_listdir(path: str) -> list[str]:
            call_count["n"] += 1
            if call_count["n"] >= 2:
                raise OSError("simulated listdir failure")
            return orig_listdir(path)

        with patch("os.listdir", side_effect=_failing_listdir):
            context.rollback_result = context.mgr.rollback_to(context.crafted_cp)
    else:
        context.rollback_result = context.mgr.rollback_to(context.crafted_cp)


@then("the rollback result should be false")
def step_assert_rollback_false(context: Context) -> None:
    assert context.rollback_result is False, "Expected rollback to return False"


@given("a SandboxCheckpoint with a non-existent snapshot_path")
def step_cp_nonexistent_snapshot(context: Context) -> None:
    context.crafted_cp = SandboxCheckpoint(
        checkpoint_id="cp-nosnap",
        sandbox_id="sb-rb2",
        plan_id="plan-rb2",
        phase="pre_execute",
        created_at=datetime.now(),
        metadata={},
        snapshot_path="/nonexistent/snapshot/dir",
    )


@given("a SandboxCheckpoint whose snapshot exists but metadata has no sandbox_path")
def step_cp_no_sandbox_path_meta(context: Context) -> None:
    # Create a real snapshot dir so the first guard passes
    snap_parent = tempfile.mkdtemp(prefix="ca-rb-test-")
    snap = os.path.join(snap_parent, "snapshot")
    os.makedirs(snap, exist_ok=True)
    _register_cleanup(context, snap_parent)
    context.crafted_cp = SandboxCheckpoint(
        checkpoint_id="cp-nometa",
        sandbox_id="sb-rb3",
        plan_id="plan-rb3",
        phase="pre_execute",
        created_at=datetime.now(),
        metadata={},  # no sandbox_path key
        snapshot_path=snap,
    )


@given("a SandboxCheckpoint whose snapshot exists but sandbox_path points nowhere")
def step_cp_sandbox_path_missing(context: Context) -> None:
    snap_parent = tempfile.mkdtemp(prefix="ca-rb-test2-")
    snap = os.path.join(snap_parent, "snapshot")
    os.makedirs(snap, exist_ok=True)
    _register_cleanup(context, snap_parent)
    context.crafted_cp = SandboxCheckpoint(
        checkpoint_id="cp-badpath",
        sandbox_id="sb-rb4",
        plan_id="plan-rb4",
        phase="pre_execute",
        created_at=datetime.now(),
        metadata={"sandbox_path": "/nonexistent/sandbox/dir"},
        snapshot_path=snap,
    )


@given("a checkpoint was created for the temporary sandbox")
def step_create_cp_for_temp_sandbox(context: Context) -> None:
    context.cp_result = context.mgr.create_checkpoint(
        sandbox=context.mock_sb,
        plan_id="plan-rb-ok",
        phase="pre_execute",
        metadata={"sandbox_path": context.temp_sandbox},
    )
    _register_cleanup(context, os.path.dirname(context.cp_result.snapshot_path))
    context.crafted_cp = context.cp_result


@given("the temporary sandbox files are then modified")
def step_modify_temp_sandbox(context: Context) -> None:
    # Overwrite existing file
    with open(os.path.join(context.temp_sandbox, "file.txt"), "w") as f:
        f.write("modified!\n")
    # Add a new file
    with open(os.path.join(context.temp_sandbox, "extra.txt"), "w") as f:
        f.write("extra\n")
    # Add a new subdirectory with file
    new_sub = os.path.join(context.temp_sandbox, "newdir")
    os.makedirs(new_sub, exist_ok=True)
    with open(os.path.join(new_sub, "added.txt"), "w") as f:
        f.write("added\n")


@when("I invoke rollback_to on the created checkpoint")
def step_rollback_created_cp(context: Context) -> None:
    context.rollback_result = context.mgr.rollback_to(context.cp_result)


@then("the rollback result should be true")
def step_assert_rollback_true(context: Context) -> None:
    assert context.rollback_result is True, "Expected rollback to return True"


@then("the temporary sandbox should contain the original files")
def step_assert_temp_sandbox_restored(context: Context) -> None:
    with open(os.path.join(context.temp_sandbox, "file.txt")) as f:
        assert f.read() == "original\n", "file.txt should be restored"
    with open(os.path.join(context.temp_sandbox, "subdir", "nested.txt")) as f:
        assert f.read() == "nested-original\n", "nested.txt should be restored"
    assert not os.path.exists(os.path.join(context.temp_sandbox, "extra.txt")), (
        "extra.txt should not exist after rollback"
    )
    assert not os.path.exists(os.path.join(context.temp_sandbox, "newdir")), (
        "newdir should not exist after rollback"
    )


@given("a SandboxCheckpoint whose snapshot and sandbox exist but restore will fail")
def step_cp_rollback_oserror(context: Context) -> None:
    # Create a real snapshot dir with a file
    snap_parent = tempfile.mkdtemp(prefix="ca-rb-err-")
    snap = os.path.join(snap_parent, "snapshot")
    os.makedirs(snap, exist_ok=True)
    with open(os.path.join(snap, "ok.txt"), "w") as f:
        f.write("ok\n")

    # Create a real sandbox dir
    sandbox_dir = tempfile.mkdtemp(prefix="ca-rb-err-sb-")
    with open(os.path.join(sandbox_dir, "existing.txt"), "w") as f:
        f.write("existing\n")

    _register_cleanup(context, snap_parent)
    _register_cleanup(context, sandbox_dir)

    context.crafted_cp = SandboxCheckpoint(
        checkpoint_id="cp-oserr",
        sandbox_id="sb-rb-err",
        plan_id="plan-rb-err",
        phase="pre_execute",
        created_at=datetime.now(),
        metadata={"sandbox_path": sandbox_dir},
        snapshot_path=snap,
    )
    # Flag the OSError scenario so the when step knows to patch
    context.rollback_should_patch_oserror = True


# ===========================================================================
# list_checkpoints - branches
# ===========================================================================


@when('I invoke list_checkpoints for sandbox "{sandbox_id}"')
def step_list_checkpoints_by_id(context: Context, sandbox_id: str) -> None:
    context.cp_list = context.mgr.list_checkpoints(sandbox_id)


@then("the checkpoint list should be empty")
def step_assert_list_empty(context: Context) -> None:
    assert context.cp_list == [], f"Expected empty list, got {context.cp_list}"


@given("I create two checkpoints for the mock sandbox")
def step_create_two_checkpoints(context: Context) -> None:
    context.mgr.create_checkpoint(
        sandbox=context.mock_sb,
        plan_id="plan-list",
        phase="pre_execute",
    )
    context.mgr.create_checkpoint(
        sandbox=context.mock_sb,
        plan_id="plan-list",
        phase="post_execute",
    )


@when("I invoke list_checkpoints for the mock sandbox")
def step_list_checkpoints_mock_sb(context: Context) -> None:
    context.cp_list = context.mgr.list_checkpoints(context.mock_sb.sandbox_id)


@then("the checkpoint list should have {count:d} entries in creation order")
def step_assert_list_count_ordered(context: Context, count: int) -> None:
    assert len(context.cp_list) == count, (
        f"Expected {count} entries, got {len(context.cp_list)}"
    )
    for i in range(1, len(context.cp_list)):
        assert context.cp_list[i].created_at >= context.cp_list[i - 1].created_at


# ===========================================================================
# delete_checkpoint - branches
# ===========================================================================


@given("a single checkpoint is created for the mock sandbox")
def step_create_single_cp(context: Context) -> None:
    context.cp_result = context.mgr.create_checkpoint(
        sandbox=context.mock_sb,
        plan_id="plan-del",
        phase="pre_execute",
    )
    _register_cleanup(context, os.path.dirname(context.cp_result.snapshot_path))


@when("I invoke delete_checkpoint with the created checkpoint id")
def step_delete_created_cp(context: Context) -> None:
    context.delete_result = context.mgr.delete_checkpoint(
        context.cp_result.checkpoint_id
    )


@then("the checkpoint delete result should be true")
def step_assert_delete_true(context: Context) -> None:
    assert context.delete_result is True, "Expected delete to return True"


@then("listing checkpoints for the mock sandbox should return {count:d} entries")
def step_assert_list_after_delete(context: Context, count: int) -> None:
    result = context.mgr.list_checkpoints(context.mock_sb.sandbox_id)
    assert len(result) == count, f"Expected {count} entries, got {len(result)}"


@when('I invoke delete_checkpoint with id "{cp_id}"')
def step_delete_unknown_cp(context: Context, cp_id: str) -> None:
    context.delete_result = context.mgr.delete_checkpoint(cp_id)


@then("the checkpoint delete result should be false")
def step_assert_delete_false(context: Context) -> None:
    assert context.delete_result is False, "Expected delete to return False"