cleveragents-core/features/steps/plan_cli_rollback_steps.py

"""Step definitions for plan CLI rollback command feature."""

from __future__ import annotations

import json
from datetime import UTC, datetime
from unittest.mock import MagicMock, patch

from behave import given, then, when
from behave.runner import Context
from typer.testing import CliRunner
import yaml

from cleveragents.cli.commands.plan import app as plan_app
from cleveragents.core.exceptions import (
    BusinessRuleViolation,
    CleverAgentsError,
    ResourceNotFoundError,
)
from cleveragents.domain.models.core.checkpoint import Checkpoint, CheckpointMetadata

_PATCH_CONTAINER = "cleveragents.application.container.get_container"
_RUNNER = CliRunner()


def _make_test_checkpoint(
    checkpoint_id: str,
    plan_id: str,
    index: int = 0,
) -> Checkpoint:
    """Build a Checkpoint domain object for testing."""
    return Checkpoint(
        checkpoint_id=checkpoint_id,
        plan_id=plan_id,
        sandbox_ref="test-sandbox",
        checkpoint_type="manual",
        created_at=datetime(2026, 4, 15, 0, index, 0, tzinfo=UTC),
        metadata=CheckpointMetadata(
            reason=f"checkpoint-{index}",
            source_tool="test",
            phase="execute",
        ),
    )


def _build_mock_container(context: Context) -> MagicMock:
    """Build a MagicMock container from context state."""
    mock_container = MagicMock()
    mock_cp_svc = MagicMock()
    mock_lifecycle_svc = MagicMock()

    plan_id = getattr(context, "plan_id", None)
    checkpoint_id = getattr(context, "checkpoint_id", None)

    if plan_id and hasattr(context, "plans") and plan_id in context.plans:
        plan_data = context.plans[plan_id]
        raw_checkpoints = plan_data.get("checkpoints", [])
        cp_objects = [
            _make_test_checkpoint(
                checkpoint_id=cp["id"],
                plan_id=plan_id,
                index=i,
            )
            for i, cp in enumerate(raw_checkpoints)
        ]
        mock_cp_svc.list_checkpoints.return_value = cp_objects
        # Use fallback prompt text for confirmation flows
        mock_cp_svc.get_checkpoint.side_effect = CleverAgentsError(
            "metadata unavailable"
        )

        if plan_data.get("applied", False):
            mock_lifecycle_svc.rollback_plan.side_effect = BusinessRuleViolation(
                "plan is already applied"
            )
        else:
            mock_result = MagicMock()
            mock_result.from_checkpoint_id = checkpoint_id or ""
            mock_result.restored_files_count = len(raw_checkpoints)
            mock_result.changed_paths = [
                f"file{i}.py" for i in range(len(raw_checkpoints))
            ]
            mock_result.changes_reverted = None
            mock_lifecycle_svc.rollback_plan.return_value = mock_result

            if checkpoint_id:
                matching = [
                    cp for cp in cp_objects if cp.checkpoint_id == checkpoint_id
                ]
                if not matching:
                    mock_lifecycle_svc.rollback_plan.side_effect = (
                        ResourceNotFoundError(
                            resource_type="checkpoint", resource_id=checkpoint_id
                        )
                    )
    elif plan_id == "INVALID_PLAN_ID":
        mock_cp_svc.list_checkpoints.side_effect = CleverAgentsError("invalid plan ID")
        mock_lifecycle_svc.rollback_plan.side_effect = CleverAgentsError(
            "invalid plan ID"
        )
    else:
        mock_cp_svc.list_checkpoints.side_effect = ResourceNotFoundError(
            resource_type="plan", resource_id=plan_id or "unknown"
        )
        mock_lifecycle_svc.rollback_plan.side_effect = ResourceNotFoundError(
            resource_type="plan", resource_id=plan_id or "unknown"
        )

    mock_container.checkpoint_service.return_value = mock_cp_svc
    mock_container.plan_lifecycle_service.return_value = mock_lifecycle_svc
    return mock_container


def _invoke_rollback(
    context: Context,
    args: list[str],
    input_text: str | None = None,
) -> None:
    """Invoke the rollback CLI command with a mocked container."""
    mock_container = _build_mock_container(context)
    with patch(_PATCH_CONTAINER, return_value=mock_container):
        result = _RUNNER.invoke(plan_app, args, input=input_text)
    context.last_exit_code = result.exit_code
    context.last_stdout = result.output
    context.last_stderr = str(result.exception) if result.exception else ""


# ----------------------------------------------------------------------
# Given steps
# ----------------------------------------------------------------------


@given('a plan "{plan_id}" exists')
def step_plan_exists(context: Context, plan_id: str) -> None:
    """Create a plan with the given ID."""
    context.plan_id = plan_id
    if not hasattr(context, "plans"):
        context.plans = {}
    context.plans[plan_id] = {
        "id": plan_id,
        "checkpoints": [],
        "applied": False,
    }


@given("the plan has {count:d} checkpoints")
def step_plan_has_checkpoints(context: Context, count: int) -> None:
    """Add N auto-generated checkpoints to the plan."""
    plan_id = context.plan_id
    if plan_id not in context.plans:
        context.plans[plan_id] = {"id": plan_id, "checkpoints": [], "applied": False}
    for i in range(count):
        checkpoint_id = f"01BRZ4NFEKTSV4RRFFQ69G5F{i:02d}"
        context.plans[plan_id]["checkpoints"].append(
            {
                "id": checkpoint_id,
                "created_at": f"2026-04-15T00:{i:02d}:00Z",
                "reason": f"checkpoint-{i}",
            }
        )


@given("the plan has no checkpoints")
def step_plan_has_no_checkpoints(context: Context) -> None:
    """Ensure the plan has no checkpoints."""
    plan_id = context.plan_id
    if plan_id not in context.plans:
        context.plans[plan_id] = {"id": plan_id, "checkpoints": [], "applied": False}
    context.plans[plan_id]["checkpoints"] = []


@given('the plan has a checkpoint "{checkpoint_id}"')
def step_plan_has_checkpoint(context: Context, checkpoint_id: str) -> None:
    """Add a specific checkpoint to the plan."""
    plan_id = context.plan_id
    if plan_id not in context.plans:
        context.plans[plan_id] = {"id": plan_id, "checkpoints": [], "applied": False}
    context.plans[plan_id]["checkpoints"].append(
        {
            "id": checkpoint_id,
            "created_at": "2026-04-15T00:00:00Z",
            "reason": "test-checkpoint",
        }
    )
    context.checkpoint_id = checkpoint_id


@given("the plan is not applied")
def step_plan_not_applied(context: Context) -> None:
    """Mark the plan as not applied."""
    plan_id = context.plan_id
    if plan_id not in context.plans:
        context.plans[plan_id] = {"id": plan_id, "checkpoints": [], "applied": False}
    context.plans[plan_id]["applied"] = False


@given("the plan is marked as applied")
def step_plan_marked_applied(context: Context) -> None:
    """Mark the plan as applied."""
    plan_id = context.plan_id
    if plan_id not in context.plans:
        context.plans[plan_id] = {"id": plan_id, "checkpoints": [], "applied": False}
    context.plans[plan_id]["applied"] = True


# ----------------------------------------------------------------------
# When steps
# ----------------------------------------------------------------------


@when('I run "agents plan rollback {plan_id:S}"')
def step_run_rollback_list(context: Context, plan_id: str) -> None:
    """Run the rollback command in list mode (no checkpoint ID)."""
    context.plan_id = plan_id
    _invoke_rollback(context, ["rollback", plan_id])


@when('I run "agents plan rollback {plan_id:S} {checkpoint_id:S}"')
def step_run_rollback_restore(
    context: Context, plan_id: str, checkpoint_id: str
) -> None:
    """Run the rollback command in restore mode (no flags)."""
    context.plan_id = plan_id
    context.checkpoint_id = checkpoint_id
    _invoke_rollback(context, ["rollback", plan_id, checkpoint_id])


@when('I run "agents plan rollback {plan_id:S} {checkpoint_id:S} --yes"')
def step_run_rollback_restore_yes(
    context: Context, plan_id: str, checkpoint_id: str
) -> None:
    """Run the rollback command with --yes to skip confirmation."""
    context.plan_id = plan_id
    context.checkpoint_id = checkpoint_id
    _invoke_rollback(context, ["rollback", plan_id, checkpoint_id, "--yes"])


@when('I run "agents plan rollback {plan_id:S} {checkpoint_id:S}" with confirmation')
def step_run_rollback_with_confirmation(
    context: Context, plan_id: str, checkpoint_id: str
) -> None:
    """Run the rollback command and confirm the prompt."""
    context.plan_id = plan_id
    context.checkpoint_id = checkpoint_id
    _invoke_rollback(context, ["rollback", plan_id, checkpoint_id], input_text="y\n")


@when(
    'I run "agents plan rollback {plan_id:S} {checkpoint_id:S}" and cancel confirmation'
)
def step_run_rollback_cancel_confirmation(
    context: Context, plan_id: str, checkpoint_id: str
) -> None:
    """Run the rollback command and cancel the confirmation prompt."""
    context.plan_id = plan_id
    context.checkpoint_id = checkpoint_id
    _invoke_rollback(context, ["rollback", plan_id, checkpoint_id], input_text="n\n")


@when('I run "agents plan rollback {plan_id:S} --format {fmt:S}"')
def step_run_rollback_with_format(context: Context, plan_id: str, fmt: str) -> None:
    """Run the rollback (list mode) with a specific output format."""
    context.plan_id = plan_id
    context.format = fmt
    _invoke_rollback(context, ["rollback", plan_id, "--format", fmt])


@when(
    'I run "agents plan rollback {plan_id:S} {checkpoint_id:S} --yes --format {fmt:S}"'
)
def step_run_rollback_restore_with_format(
    context: Context, plan_id: str, checkpoint_id: str, fmt: str
) -> None:
    """Run the rollback command in restore mode with a specific format."""
    context.plan_id = plan_id
    context.checkpoint_id = checkpoint_id
    context.format = fmt
    _invoke_rollback(
        context, ["rollback", plan_id, checkpoint_id, "--yes", "--format", fmt]
    )


@when(
    'I run "agents plan rollback {plan_id:S} --to-checkpoint {checkpoint_id:S} --yes"'
)
def step_run_rollback_to_checkpoint(
    context: Context, plan_id: str, checkpoint_id: str
) -> None:
    """Run the rollback command using the --to-checkpoint named option."""
    context.plan_id = plan_id
    context.checkpoint_id = checkpoint_id
    _invoke_rollback(
        context, ["rollback", plan_id, "--to-checkpoint", checkpoint_id, "--yes"]
    )


# ----------------------------------------------------------------------
# Then steps
# ----------------------------------------------------------------------


@then("the output should list {count:d} checkpoints")
def step_output_lists_checkpoints(context: Context, count: int) -> None:
    """Verify the output lists the expected number of checkpoints."""
    output = context.last_stdout
    checkpoint_count = output.count("01BRZ4NFEKTSV4RRFFQ69G5F")
    assert checkpoint_count == count, (
        f"Expected {count} checkpoints, found {checkpoint_count}"
    )


@then("each checkpoint should show ID, timestamp, and metadata")
def step_checkpoint_shows_metadata(context: Context) -> None:
    """Verify each checkpoint shows required metadata."""
    output = context.last_stdout
    assert "01BRZ4NFEKTSV4RRFFQ69G5F" in output, "Checkpoint ID not found in output"
    assert "2026-04-15" in output or "checkpoint" in output.lower(), (
        "Checkpoint metadata not found"
    )


@then("the output should indicate no checkpoints available")
def step_output_no_checkpoints(context: Context) -> None:
    """Verify the output indicates no checkpoints are available."""
    output = context.last_stdout.lower()
    assert any(
        phrase in output for phrase in ["no checkpoints", "empty", "0 checkpoints"]
    ), f"Output does not indicate no checkpoints: {context.last_stdout}"


@then("the output should be formatted as a rich table")
def step_output_rich_table(context: Context) -> None:
    """Verify the output is non-empty (rich table rendered)."""
    assert len(context.last_stdout) > 0, "Output is empty"


@then("the table should contain checkpoint IDs and metadata")
def step_table_contains_metadata(context: Context) -> None:
    """Verify the table contains checkpoint IDs."""
    assert "01BRZ4NFEKTSV4RRFFQ69G5F" in context.last_stdout, (
        "Checkpoint ID not in table"
    )


@then("the rollback output should be valid JSON")
def step_output_valid_json(context: Context) -> None:
    """Verify the output is valid JSON."""
    try:
        context.json_output = json.loads(context.last_stdout)
    except json.JSONDecodeError as e:
        raise AssertionError(f"Output is not valid JSON: {e}") from e


@then("the JSON should contain a list of checkpoints")
def step_json_contains_checkpoints(context: Context) -> None:
    """Verify the JSON contains a checkpoints list."""
    assert isinstance(context.json_output, dict), "JSON output should be a dict"
    inner = context.json_output.get("data", context.json_output)
    assert "checkpoints" in inner, "JSON should contain 'checkpoints' key"
    assert isinstance(inner["checkpoints"], list), "Checkpoints should be a list"


@then("the rollback output should be valid YAML")
def step_output_valid_yaml(context: Context) -> None:
    """Verify the output is valid YAML."""
    try:
        context.yaml_output = yaml.safe_load(context.last_stdout)
    except yaml.YAMLError as e:
        raise AssertionError(f"Output is not valid YAML: {e}") from e


@then("the YAML should contain a list of checkpoints")
def step_yaml_contains_checkpoints(context: Context) -> None:
    """Verify the YAML contains a checkpoints list."""
    assert isinstance(context.yaml_output, dict), "YAML output should be a dict"
    inner = context.yaml_output.get("data", context.yaml_output)
    assert "checkpoints" in inner, "YAML should contain 'checkpoints' key"
    assert isinstance(inner["checkpoints"], list), "Checkpoints should be a list"


@then("the output should be plain text")
def step_output_plain_text(context: Context) -> None:
    """Verify the output is non-empty plain text."""
    assert len(context.last_stdout) > 0, "Output is empty"


@then("the text should contain checkpoint IDs")
def step_text_contains_checkpoint_ids(context: Context) -> None:
    """Verify the text contains checkpoint IDs."""
    assert "01BRZ4NFEKTSV4RRFFQ69G5F" in context.last_stdout, (
        "Checkpoint ID not in output"
    )


@then("the plan rollback should succeed")
def step_rollback_succeeds(context: Context) -> None:
    """Verify the rollback succeeded (exit code 0)."""
    assert context.last_exit_code == 0, (
        f"Rollback failed with exit code {context.last_exit_code}: {context.last_stdout}"
    )


@then("the output should show rollback summary")
def step_output_shows_summary(context: Context) -> None:
    """Verify the output shows a rollback summary."""
    output = context.last_stdout.lower()
    assert any(phrase in output for phrase in ["rollback", "restored", "summary"]), (
        f"Output does not show rollback summary: {context.last_stdout}"
    )


@then("no rollback confirmation prompt should be shown")
def step_no_confirmation_prompt(context: Context) -> None:
    """Verify the --yes flag bypassed the confirmation prompt."""
    assert context.last_exit_code == 0, "Command should succeed with --yes flag"


@then("the command should fail with not found error")
def step_command_fails_not_found(context: Context) -> None:
    """Verify the command fails with a not-found error."""
    assert context.last_exit_code != 0, "Command should fail"
    error_output = (context.last_stdout + context.last_stderr).lower()
    assert any(phrase in error_output for phrase in ["not found", "does not exist"]), (
        f"Error message does not indicate not found: {context.last_stdout}"
    )


@then('the command should fail with "{error_message}"')
def step_command_fails_with_message(context: Context, error_message: str) -> None:
    """Verify the command fails with a specific error message."""
    assert context.last_exit_code != 0, "Command should fail"
    error_output = (context.last_stdout + context.last_stderr).lower()
    assert error_message.lower() in error_output, (
        f"Error message '{error_message}' not found in: {context.last_stdout}"
    )


@then("the JSON should contain rollback summary")
def step_json_contains_rollback_summary(context: Context) -> None:
    """Verify the JSON contains a rollback_summary key."""
    assert isinstance(context.json_output, dict), "JSON output should be a dict"
    inner = context.json_output.get("data", context.json_output)
    assert "rollback_summary" in inner, "JSON should contain 'rollback_summary' key"


@then("the YAML should contain rollback summary")
def step_yaml_contains_rollback_summary(context: Context) -> None:
    """Verify the YAML contains a rollback_summary key."""
    assert isinstance(context.yaml_output, dict), "YAML output should be a dict"
    inner = context.yaml_output.get("data", context.yaml_output)
    assert "rollback_summary" in inner, "YAML should contain 'rollback_summary' key"


@then("the rollback should be atomic")
def step_rollback_atomic(context: Context) -> None:
    """Verify the rollback completed atomically (exit code 0)."""
    assert context.last_exit_code == 0, "Rollback should succeed atomically"


@then("the plan state should be consistent")
def step_plan_state_consistent(context: Context) -> None:
    """Verify the plan state is consistent after rollback."""
    assert context.last_exit_code == 0, "Plan state should be consistent after rollback"


@then("the output should show the number of restored files")
def step_output_shows_restored_files(context: Context) -> None:
    """Verify the output mentions restored files."""
    output = context.last_stdout.lower()
    assert any(phrase in output for phrase in ["restored", "files", "changes"]), (
        f"Output does not show restored files: {context.last_stdout}"
    )


@then("the output should list changed paths")
def step_output_lists_changed_paths(context: Context) -> None:
    """Verify the output lists changed paths."""
    output = context.last_stdout.lower()
    assert any(phrase in output for phrase in ["changed", "paths", "files"]), (
        f"Output does not list changed paths: {context.last_stdout}"
    )


@then("the command should fail with validation error")
def step_command_fails_validation(context: Context) -> None:
    """Verify the command fails with a validation error."""
    assert context.last_exit_code != 0, "Command should fail"
    error_output = (context.last_stdout + context.last_stderr).lower()
    assert any(
        phrase in error_output for phrase in ["invalid", "validation", "error"]
    ), f"Error message does not indicate validation error: {context.last_stdout}"


@then("the output should list available checkpoints")
def step_output_lists_available_checkpoints(context: Context) -> None:
    """Verify the output lists available checkpoints (list mode)."""
    output = context.last_stdout.lower()
    assert any(phrase in output for phrase in ["checkpoint", "no checkpoints"]), (
        f"Output does not list checkpoints: {context.last_stdout}"
    )


@then("the rollback should be cancelled")
def step_rollback_cancelled(context: Context) -> None:
    """Verify the rollback was cancelled by the user."""
    assert context.last_exit_code != 0, "Rollback should be cancelled"
    output = context.last_stdout.lower()
    assert any(phrase in output for phrase in ["cancelled", "aborted"]), (
        f"Output does not indicate rollback was cancelled: {context.last_stdout}"
    )


@then("the plan state should remain unchanged")
def step_plan_state_unchanged(context: Context) -> None:
    """Verify the plan state did not change (rollback was not completed)."""
    assert context.last_exit_code != 0, "Rollback should not have completed"