cleveragents-core/features/steps/correction_model_steps.py

"""Step definitions for correction_model.feature."""

from __future__ import annotations

from behave import given, then, when
from pydantic import ValidationError as PydanticValidationError

from cleveragents.application.services.correction_service import CorrectionService
from cleveragents.core.exceptions import ResourceNotFoundError, ValidationError
from cleveragents.domain.models.core.correction import (
    CorrectionAttempt,
    CorrectionImpact,
    CorrectionMode,
    CorrectionRequest,
    CorrectionResult,
    CorrectionStatus,
)

# ── Given steps ──────────────────────────────────────────────────


@given("a correction service")
def step_given_correction_service(context):
    context.correction_service = CorrectionService()
    context.correction = None
    context.impact = None
    context.result = None
    context.error = None


@given(
    'a correction request for plan "{plan_id}" decision "{decision_id}" '
    'in {mode} mode with guidance "{guidance}"'
)
def step_given_correction_request(context, plan_id, decision_id, mode, guidance):
    svc = context.correction_service
    correction_mode = CorrectionMode(mode)
    context.correction = svc.request_correction(
        plan_id=plan_id,
        target_decision_id=decision_id,
        mode=correction_mode,
        guidance=guidance,
    )


@given('a cancelled correction for plan "{plan_id}" decision "{decision_id}"')
def step_given_cancelled_correction(context, plan_id, decision_id):
    svc = context.correction_service
    correction = svc.request_correction(
        plan_id=plan_id,
        target_decision_id=decision_id,
        mode=CorrectionMode.REVERT,
        guidance="To be cancelled",
    )
    svc.cancel_correction(correction.correction_id)
    context.correction = svc.get_correction(correction.correction_id)


@given('an applied correction for plan "{plan_id}" decision "{decision_id}"')
def step_given_applied_correction(context, plan_id, decision_id):
    svc = context.correction_service
    correction = svc.request_correction(
        plan_id=plan_id,
        target_decision_id=decision_id,
        mode=CorrectionMode.REVERT,
        guidance="To be applied",
    )
    svc.execute_correction(correction.correction_id)
    context.correction = svc.get_correction(correction.correction_id)


# ── When steps ───────────────────────────────────────────────────


@when(
    'I request a correction for plan "{plan_id}" decision "{decision_id}" '
    'in {mode} mode with guidance "{guidance}"'
)
def step_when_request_correction(context, plan_id, decision_id, mode, guidance):
    svc = context.correction_service
    correction_mode = CorrectionMode(mode)
    context.correction = svc.request_correction(
        plan_id=plan_id,
        target_decision_id=decision_id,
        mode=correction_mode,
        guidance=guidance,
    )


@when(
    'I request a dry-run correction for plan "{plan_id}" decision '
    '"{decision_id}" in {mode} mode with guidance "{guidance}"'
)
def step_when_request_dry_run(context, plan_id, decision_id, mode, guidance):
    svc = context.correction_service
    correction_mode = CorrectionMode(mode)
    context.correction = svc.request_correction(
        plan_id=plan_id,
        target_decision_id=decision_id,
        mode=correction_mode,
        guidance=guidance,
        dry_run=True,
    )


@when("I analyze the correction impact")
def step_when_analyze_impact(context):
    svc = context.correction_service
    context.impact = svc.analyze_impact(context.correction.correction_id)


@when("I cancel the correction")
def step_when_cancel(context):
    svc = context.correction_service
    context.correction = svc.cancel_correction(context.correction.correction_id)


@when("I try to cancel the correction")
def step_when_try_cancel(context):
    svc = context.correction_service
    context.error = None
    try:
        svc.cancel_correction(context.correction.correction_id)
    except (ValidationError, ResourceNotFoundError) as exc:
        context.error = exc


@when("I execute the correction")
def step_when_execute(context):
    svc = context.correction_service
    context.result = svc.execute_correction(context.correction.correction_id)
    # Refresh the correction object
    context.correction = svc.get_correction(context.correction.correction_id)


@when("I try to execute the correction")
def step_when_try_execute(context):
    svc = context.correction_service
    context.error = None
    try:
        svc.execute_correction(context.correction.correction_id)
    except (ValidationError, ResourceNotFoundError) as exc:
        context.error = exc


@when("I get the correction by its ID")
def step_when_get_by_id(context):
    svc = context.correction_service
    context.found_correction = svc.get_correction(context.correction.correction_id)


@when('I try to get correction "{correction_id}"')
def step_when_try_get_nonexistent(context, correction_id):
    svc = context.correction_service
    context.error = None
    try:
        svc.get_correction(correction_id)
    except ResourceNotFoundError as exc:
        context.error = exc


@when('I list corrections for plan "{plan_id}"')
def step_when_list_by_plan(context, plan_id):
    svc = context.correction_service
    context.listed_corrections = svc.list_corrections(plan_id=plan_id)


@when("I list all corrections")
def step_when_list_all(context):
    svc = context.correction_service
    context.listed_corrections = svc.list_corrections()


@when("I try to request a correction with empty decision ID")
def step_when_empty_decision(context):
    svc = context.correction_service
    context.error = None
    try:
        svc.request_correction(
            plan_id="plan-1",
            target_decision_id="",
            mode=CorrectionMode.REVERT,
            guidance="test",
        )
    except ValidationError as exc:
        context.error = exc


@when("I try to request a correction with empty plan ID")
def step_when_empty_plan(context):
    svc = context.correction_service
    context.error = None
    try:
        svc.request_correction(
            plan_id="",
            target_decision_id="DEC-001",
            mode=CorrectionMode.REVERT,
            guidance="test",
        )
    except ValidationError as exc:
        context.error = exc


@when("I try to request a correction with empty guidance")
def step_when_empty_guidance(context):
    svc = context.correction_service
    context.error = None
    context.correction = svc.request_correction(
        plan_id="plan-1",
        target_decision_id="DEC-001",
        mode=CorrectionMode.REVERT,
        guidance="",
    )


@when("I try to request a correction with whitespace-only guidance")
def step_when_whitespace_guidance(context):
    svc = context.correction_service
    context.error = None
    context.correction = svc.request_correction(
        plan_id="plan-1",
        target_decision_id="DEC-001",
        mode=CorrectionMode.REVERT,
        guidance="   ",
    )


@when("I try to create a CorrectionRequest with empty plan_id")
def step_when_model_empty_plan(context):
    context.error = None
    try:
        CorrectionRequest(
            plan_id="",
            target_decision_id="DEC-001",
            mode=CorrectionMode.REVERT,
            guidance="test",
        )
    except PydanticValidationError as exc:
        context.error = exc


@when("I try to create a CorrectionRequest with empty target_decision_id")
def step_when_model_empty_decision(context):
    context.error = None
    try:
        CorrectionRequest(
            plan_id="plan-1",
            target_decision_id="",
            mode=CorrectionMode.REVERT,
            guidance="test",
        )
    except PydanticValidationError as exc:
        context.error = exc


@when("I try to create a CorrectionRequest with empty guidance")
def step_when_model_empty_guidance(context):
    context.error = None
    context.correction = CorrectionRequest(
        plan_id="plan-1",
        target_decision_id="DEC-001",
        mode=CorrectionMode.REVERT,
        guidance="",
    )


@when('I try to create a CorrectionImpact with invalid risk_level "{level}"')
def step_when_model_invalid_risk(context, level):
    context.error = None
    try:
        CorrectionImpact(risk_level=level)
    except PydanticValidationError as exc:
        context.error = exc


@when('I create a CorrectionImpact with risk_level "{level}"')
def step_when_create_impact(context, level):
    context.created_impact = CorrectionImpact(risk_level=level)


@when('I create a CorrectionResult with correction_id "{cid}" and status "{status}"')
def step_when_create_result(context, cid, status):
    context.created_result = CorrectionResult(
        correction_id=cid,
        status=CorrectionStatus(status),
    )


@when('I create a CorrectionAttempt for correction "{cid}"')
def step_when_create_attempt(context, cid):
    context.created_attempt = CorrectionAttempt(correction_id=cid)


# ── Then steps ───────────────────────────────────────────────────


@then("the correction should be created")
def step_then_created(context):
    assert context.correction is not None
    assert context.correction.correction_id


@then('the correction mode should be "{mode}"')
def step_then_mode(context, mode):
    assert context.correction.mode.value == mode


@then('the correction status should be "{status}"')
def step_then_status(context, status):
    assert context.correction.status.value == status


@then('the correction plan_id should be "{plan_id}"')
def step_then_plan_id(context, plan_id):
    assert context.correction.plan_id == plan_id


@then('the correction target_decision_id should be "{decision_id}"')
def step_then_decision_id(context, decision_id):
    assert context.correction.target_decision_id == decision_id


@then('the correction guidance should be "{guidance}"')
def step_then_guidance(context, guidance):
    assert context.correction.guidance == guidance


@then("the correction dry_run should be true")
def step_then_dry_run(context):
    assert context.correction.dry_run is True


@then("the impact should list affected decisions")
def step_then_impact_decisions(context):
    assert context.impact is not None
    assert len(context.impact.affected_decisions) > 0


@then('the impact risk level should be "{level}"')
def step_then_impact_risk(context, level):
    assert context.impact.risk_level == level


@then('a validation error should be raised with message containing "{text}"')
def step_then_validation_error(context, text):
    assert context.error is not None, "Expected an error but none was raised"
    error_msg = str(context.error)
    if hasattr(context.error, "message"):
        error_msg = context.error.message
    assert text.lower() in error_msg.lower(), (
        f"Expected '{text}' in error message, got: {error_msg}"
    )


@then("a resource not found error should be raised")
def step_then_not_found(context):
    assert context.error is not None
    assert isinstance(context.error, ResourceNotFoundError)


@then("I should get {count:d} corrections")
def step_then_count(context, count):
    assert len(context.listed_corrections) == count


@then("the correction should be found")
def step_then_found(context):
    assert context.found_correction is not None


@then("the result should have reverted decisions")
def step_then_result_reverted(context):
    assert context.result is not None
    assert len(context.result.reverted_decisions) > 0


@then("the result should have no reverted decisions")
def step_then_result_no_reverted(context):
    assert context.result is not None
    assert len(context.result.reverted_decisions) == 0


@then('the result status should be "{status}"')
def step_then_result_status(context, status):
    assert context.result.status.value == status


@then("there should be {count:d} attempt recorded")
def step_then_attempt_count(context, count):
    svc = context.correction_service
    attempts = svc.list_attempts(context.correction.correction_id)
    assert len(attempts) == count


@then("the attempt should be successful")
def step_then_attempt_success(context):
    svc = context.correction_service
    attempts = svc.list_attempts(context.correction.correction_id)
    assert attempts[-1].success is True


@then("a pydantic validation error should be raised")
def step_then_pydantic_error(context):
    assert context.error is not None
    assert isinstance(context.error, PydanticValidationError)


@then("the impact should be created")
def step_then_impact_created(context):
    assert context.created_impact is not None


@then('CorrectionMode.REVERT should equal "{value}"')
def step_then_mode_revert(context, value):
    assert CorrectionMode.REVERT.value == value


@then('CorrectionMode.APPEND should equal "{value}"')
def step_then_mode_append(context, value):
    assert CorrectionMode.APPEND.value == value


@then('CorrectionStatus.PENDING should equal "{value}"')
def step_then_status_pending(context, value):
    assert CorrectionStatus.PENDING.value == value


@then('CorrectionStatus.ANALYZING should equal "{value}"')
def step_then_status_analyzing(context, value):
    assert CorrectionStatus.ANALYZING.value == value


@then('CorrectionStatus.EXECUTING should equal "{value}"')
def step_then_status_executing(context, value):
    assert CorrectionStatus.EXECUTING.value == value


@then('CorrectionStatus.APPLIED should equal "{value}"')
def step_then_status_applied(context, value):
    assert CorrectionStatus.APPLIED.value == value


@then('CorrectionStatus.FAILED should equal "{value}"')
def step_then_status_failed(context, value):
    assert CorrectionStatus.FAILED.value == value


@then('CorrectionStatus.CANCELLED should equal "{value}"')
def step_then_status_cancelled(context, value):
    assert CorrectionStatus.CANCELLED.value == value


@then("the result should have empty new_decisions")
def step_then_empty_new(context):
    assert context.created_result.new_decisions == []


@then("the result should have empty reverted_decisions")
def step_then_empty_reverted(context):
    assert context.created_result.reverted_decisions == []


@then("the result error_message should be none")
def step_then_error_msg_none(context):
    assert context.created_result.error_message is None


@then("the attempt success should be false")
def step_then_attempt_false(context):
    assert context.created_attempt.success is False


@then("the attempt completed_at should be none")
def step_then_attempt_completed_none(context):
    assert context.created_attempt.completed_at is None


@then("the attempt details should be empty")
def step_then_attempt_details_empty(context):
    assert context.created_attempt.details == {}