cleveragents-core/robot/helper_plan_correction_diff.py

"""Helper utilities for plan correction diff Robot integration tests.

Covers:
- correction_diff() output in rich, plain, JSON, and YAML formats
- PlanError raised when correction attempt not found
"""

from __future__ import annotations

import json
import sys
from collections.abc import Callable
from datetime import UTC, datetime
from unittest.mock import MagicMock

import yaml

from cleveragents.application.services.plan_apply_service import (
    PlanApplyService,
)
from cleveragents.application.services.plan_lifecycle_service import (
    PlanLifecycleService,
)
from cleveragents.config.settings import Settings
from cleveragents.core.exceptions import PlanError, ValidationError
from cleveragents.domain.models.core.correction import (
    CorrectionAttemptRecord,
    CorrectionAttemptState,
    CorrectionMode,
)
from cleveragents.infrastructure.database.repositories import (
    CorrectionAttemptNotFoundError,
)


def _make_attempt(
    correction_id: str = "01HXM9B7Z3Q1CORR001AAAAA",
    plan_id: str = "01HXM8C2ZK4Q7C2B3F2R4PLAN",
    original_decision_id: str = "01HXM9A1C2Q7W3R5ORIG0001",
    new_decision_id: str | None = "01HXM9A1C2Q7W3R5NEW00001",
    mode: CorrectionMode = CorrectionMode.REVERT,
    state: CorrectionAttemptState = CorrectionAttemptState.COMPLETE,
    guidance: str = "Use async processing instead",
    archived_artifacts_path: str | None = None,
) -> CorrectionAttemptRecord:
    """Build a minimal CorrectionAttemptRecord for testing."""
    return CorrectionAttemptRecord(
        correction_attempt_id=correction_id,
        plan_id=plan_id,
        original_decision_id=original_decision_id,
        new_decision_id=new_decision_id,
        mode=mode,
        state=state,
        guidance=guidance,
        created_at=datetime(2025, 6, 15, 10, 30, 0, tzinfo=UTC),
        completed_at=datetime(2025, 6, 15, 10, 31, 0, tzinfo=UTC),
        archived_artifacts_path=archived_artifacts_path,
    )


def _make_apply_service_with_mocked_uow(
    attempt: CorrectionAttemptRecord,
) -> tuple[PlanApplyService, str]:
    """Create a PlanApplyService with a mocked UoW returning the given attempt."""
    settings = Settings()
    lifecycle = PlanLifecycleService(settings=settings)

    # Create an action and plan so get_plan() succeeds
    lifecycle.create_action(
        name="local/robot-correction-diff-test",
        description="Robot correction diff test",
        definition_of_done="Tests pass",
        strategy_actor="local/planner",
        execution_actor="local/executor",
    )
    plan = lifecycle.use_action(action_name="local/robot-correction-diff-test")
    plan_id = plan.identity.plan_id

    # Override the attempt plan_id to match
    attempt_with_plan = attempt.model_copy(update={"plan_id": plan_id})

    # Build a mock UnitOfWork with transaction context manager
    mock_ctx = MagicMock()
    mock_ctx.correction_attempts.get.return_value = attempt_with_plan

    mock_uow = MagicMock()
    mock_uow.transaction.return_value.__enter__.return_value = mock_ctx
    mock_uow.transaction.return_value.__exit__.return_value = False

    svc = PlanApplyService(
        lifecycle_service=lifecycle,
        unit_of_work=mock_uow,
    )
    return svc, plan_id


def _correction_diff_rich() -> None:
    """Test correction_diff() rich format output."""
    attempt = _make_attempt()
    svc, plan_id = _make_apply_service_with_mocked_uow(attempt)

    output = svc.correction_diff(plan_id, attempt.correction_attempt_id, fmt="rich")
    assert "Correction Diff" in output, f"Missing 'Correction Diff': {output[:300]}"
    assert "Comparison" in output, f"Missing 'Comparison': {output[:300]}"
    assert "Patch Preview" in output, f"Missing 'Patch Preview': {output[:300]}"
    assert attempt.correction_attempt_id in output, (
        f"Correction ID missing: {output[:300]}"
    )
    assert attempt.original_decision_id in output, (
        f"Original decision missing: {output[:300]}"
    )
    print("correction-diff-rich-ok")


def _correction_diff_plain() -> None:
    """Test correction_diff() plain format output."""
    attempt = _make_attempt()
    svc, plan_id = _make_apply_service_with_mocked_uow(attempt)

    output = svc.correction_diff(plan_id, attempt.correction_attempt_id, fmt="plain")
    assert "Correction Diff" in output, f"Missing 'Correction Diff': {output[:300]}"
    assert "Comparison" in output, f"Missing 'Comparison': {output[:300]}"
    assert "Patch Preview" in output, f"Missing 'Patch Preview': {output[:300]}"
    assert attempt.correction_attempt_id in output, (
        f"Correction ID missing: {output[:300]}"
    )
    print("correction-diff-plain-ok")


def _correction_diff_json() -> None:
    """Test correction_diff() JSON format output."""
    attempt = _make_attempt()
    svc, plan_id = _make_apply_service_with_mocked_uow(attempt)

    output = svc.correction_diff(plan_id, attempt.correction_attempt_id, fmt="json")
    parsed = json.loads(output)
    assert "correction_diff" in parsed, f"Missing 'correction_diff' key: {list(parsed)}"
    assert "comparison" in parsed, f"Missing 'comparison' key: {list(parsed)}"
    assert "patch_preview" in parsed, f"Missing 'patch_preview' key: {list(parsed)}"
    assert parsed["correction_diff"]["correction"] == attempt.correction_attempt_id
    assert parsed["correction_diff"]["mode"] == "revert"
    print("correction-diff-json-ok")


def _correction_diff_yaml() -> None:
    """Test correction_diff() YAML format output."""
    attempt = _make_attempt()
    svc, plan_id = _make_apply_service_with_mocked_uow(attempt)

    output = svc.correction_diff(plan_id, attempt.correction_attempt_id, fmt="yaml")
    parsed = yaml.safe_load(output)
    assert isinstance(parsed, dict), (
        f"Expected dict, got {type(parsed)}: {output[:300]}"
    )
    assert "correction_diff" in parsed, f"Missing 'correction_diff' key: {list(parsed)}"
    assert "comparison" in parsed, f"Missing 'comparison' key: {list(parsed)}"
    assert "patch_preview" in parsed, f"Missing 'patch_preview' key: {list(parsed)}"
    assert parsed["correction_diff"]["correction"] == attempt.correction_attempt_id
    print("correction-diff-yaml-ok")


def _correction_not_found() -> None:
    """Test correction_diff() raises PlanError when correction not found."""
    settings = Settings()
    lifecycle = PlanLifecycleService(settings=settings)

    lifecycle.create_action(
        name="local/robot-correction-not-found",
        description="Robot correction not found test",
        definition_of_done="Tests pass",
        strategy_actor="local/planner",
        execution_actor="local/executor",
    )
    plan = lifecycle.use_action(action_name="local/robot-correction-not-found")
    plan_id = plan.identity.plan_id

    # Mock UoW that raises CorrectionAttemptNotFoundError
    mock_ctx = MagicMock()
    mock_ctx.correction_attempts.get.side_effect = CorrectionAttemptNotFoundError(
        "Not found"
    )
    mock_uow = MagicMock()
    mock_uow.transaction.return_value.__enter__.return_value = mock_ctx
    mock_uow.transaction.return_value.__exit__.return_value = False

    svc = PlanApplyService(lifecycle_service=lifecycle, unit_of_work=mock_uow)

    try:
        svc.correction_diff(plan_id, "CORR-NONEXISTENT", fmt="plain")
        raise AssertionError("Should have raised PlanError")
    except PlanError as e:
        assert "not found" in e.message.lower(), f"Unexpected message: {e.message}"
    print("correction-not-found-ok")


def _correction_diff_uow_none() -> None:
    """Test correction_diff() raises ValidationError when unit_of_work is None."""
    settings = Settings()
    lifecycle = PlanLifecycleService(settings=settings)

    svc = PlanApplyService(lifecycle_service=lifecycle, unit_of_work=None)

    try:
        svc.correction_diff("some-plan-id", "some-correction-id", fmt="plain")
        raise AssertionError("Should have raised ValidationError")
    except ValidationError as e:
        assert "unit_of_work is required" in e.message, (
            f"Unexpected message: {e.message}"
        )
    print("correction-diff-uow-none-ok")


COMMANDS: dict[str, Callable[[], None]] = {
    "correction-diff-rich": _correction_diff_rich,
    "correction-diff-plain": _correction_diff_plain,
    "correction-diff-json": _correction_diff_json,
    "correction-diff-yaml": _correction_diff_yaml,
    "correction-not-found": _correction_not_found,
    "correction-diff-uow-none": _correction_diff_uow_none,
}


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(f"Usage: {sys.argv[0]} <command>", file=sys.stderr)
        print(f"Available: {', '.join(COMMANDS)}", file=sys.stderr)
        sys.exit(1)
    cmd = sys.argv[1]
    func = COMMANDS.get(cmd)
    if func is None:
        print(f"Unknown command: {cmd}", file=sys.stderr)
        sys.exit(1)
    func()