cleveragents-core/robot/helper_correction_flows.py

"""Helper script for correction flows Robot Framework smoke tests.

Exercises CorrectionService revert/append flows, impact analysis,
dry-run reports, and cancellation without requiring persistence.
"""

from __future__ import annotations

import sys

from cleveragents.application.services.correction_service import CorrectionService
from cleveragents.core.exceptions import ValidationError
from cleveragents.domain.models.core.correction import CorrectionMode

# ---------------------------------------------------------------------------
# Shared fixtures
# ---------------------------------------------------------------------------
_PLAN_ID = "01ARZ3NDEKTSV4RRFFQ69G5FAV"
_DECISION_D1 = "D1"
_DECISION_D2 = "D2"
_DECISION_D3 = "D3"


def _simple_tree() -> dict[str, list[str]]:
    return {_DECISION_D1: [_DECISION_D2, _DECISION_D3]}


def _deep_tree() -> dict[str, list[str]]:
    return {
        _DECISION_D1: [_DECISION_D2, _DECISION_D3],
        _DECISION_D2: ["D4", "D5"],
        _DECISION_D3: ["D6"],
    }


# ---------------------------------------------------------------------------
# Command handlers
# ---------------------------------------------------------------------------


def _revert_simple() -> None:
    svc = CorrectionService()
    req = svc.request_correction(_PLAN_ID, _DECISION_D1, CorrectionMode.REVERT)
    result = svc.execute_revert(req.correction_id, {_DECISION_D1: []})
    assert result.status == "applied", f"Expected applied, got {result.status}"
    assert _DECISION_D1 in result.reverted_decisions
    print("revert-simple-ok")


def _revert_deep_subtree() -> None:
    svc = CorrectionService()
    req = svc.request_correction(_PLAN_ID, _DECISION_D1, CorrectionMode.REVERT)
    result = svc.execute_revert(req.correction_id, _deep_tree())
    assert len(result.reverted_decisions) == 6
    print("revert-deep-subtree-ok")


def _append_correction() -> None:
    svc = CorrectionService()
    req = svc.request_correction(_PLAN_ID, _DECISION_D1, CorrectionMode.APPEND)
    result = svc.execute_append(req.correction_id)
    assert result.status == "applied"
    assert result.spawned_child_plan_id is not None
    print("append-correction-ok")


def _dry_run_report() -> None:
    svc = CorrectionService()
    req = svc.request_correction(_PLAN_ID, _DECISION_D1, CorrectionMode.REVERT)
    report = svc.generate_dry_run_report(req.correction_id, _simple_tree())
    assert report.mode == "revert"
    assert len(report.decisions_to_invalidate) == 3
    assert report.estimated_recompute_time_seconds > 0
    print("dry-run-report-ok")


def _impact_risk_levels() -> None:
    svc = CorrectionService()
    # Low risk: 2 nodes
    req1 = svc.request_correction(_PLAN_ID, _DECISION_D1, CorrectionMode.REVERT)
    impact1 = svc.analyze_impact(req1.correction_id, {_DECISION_D1: [_DECISION_D2]})
    assert impact1.risk_level == "low", f"Expected low, got {impact1.risk_level}"

    # High risk: > 10 nodes (chain)
    tree: dict[str, list[str]] = {}
    current = "R"
    for i in range(14):
        child = f"R_c{i}"
        tree[current] = [child]
        current = child
    req2 = svc.request_correction(_PLAN_ID, "R", CorrectionMode.REVERT)
    impact2 = svc.analyze_impact(req2.correction_id, tree)
    assert impact2.risk_level == "high", f"Expected high, got {impact2.risk_level}"
    print("impact-risk-levels-ok")


def _cancel_correction() -> None:
    svc = CorrectionService()
    req = svc.request_correction(_PLAN_ID, _DECISION_D1, CorrectionMode.REVERT)
    svc.cancel_correction(req.correction_id)
    updated = svc.get_correction(req.correction_id)
    assert updated.status == "cancelled"
    print("cancel-correction-ok")


def _execute_cancelled_error() -> None:
    svc = CorrectionService()
    req = svc.request_correction(_PLAN_ID, _DECISION_D1, CorrectionMode.REVERT)
    svc.cancel_correction(req.correction_id)
    try:
        svc.execute_revert(req.correction_id, {})
        print("FAIL: expected ValidationError")
        sys.exit(1)
    except ValidationError:
        print("execute-cancelled-error-ok")


def _list_corrections_by_plan() -> None:
    svc = CorrectionService()
    svc.request_correction(_PLAN_ID, _DECISION_D1, CorrectionMode.REVERT)
    svc.request_correction(_PLAN_ID, _DECISION_D2, CorrectionMode.APPEND)
    svc.request_correction("OTHER_PLAN", _DECISION_D3, CorrectionMode.REVERT)
    result = svc.list_corrections(plan_id=_PLAN_ID)
    assert len(result) == 2, f"Expected 2, got {len(result)}"
    print("list-corrections-by-plan-ok")


# ---------------------------------------------------------------------------
# Dispatch
# ---------------------------------------------------------------------------

COMMANDS: dict[str, object] = {
    "revert-simple": _revert_simple,
    "revert-deep-subtree": _revert_deep_subtree,
    "append-correction": _append_correction,
    "dry-run-report": _dry_run_report,
    "impact-risk-levels": _impact_risk_levels,
    "cancel-correction": _cancel_correction,
    "execute-cancelled-error": _execute_cancelled_error,
    "list-corrections-by-plan": _list_corrections_by_plan,
}


def main() -> None:
    """Dispatch command from sys.argv."""
    if len(sys.argv) < 2:
        raise SystemExit("Expected command argument")
    command = sys.argv[1]
    if command not in COMMANDS:
        raise SystemExit(f"Unknown command: {command}")
    func = COMMANDS[command]
    if callable(func):
        func()


if __name__ == "__main__":
    main()