cleveragents-core/features/steps/checkpoint_manager_rollback_steps.py

"""Step definitions for rollback_to, list_checkpoints, and delete_checkpoint branches.

Split from checkpoint_manager_coverage_steps.py to keep files under 500 lines.
"""

from __future__ import annotations

import os
import shutil
import tempfile
from datetime import datetime
from unittest.mock import patch

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

from cleveragents.infrastructure.sandbox.checkpoint import SandboxCheckpoint


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 p=path: shutil.rmtree(p, ignore_errors=True)
    )


# ===========================================================================
# 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"