DataOps-platform/tests/agent/test_recovery_scenarios.py

import copy

import pytest

from app.core.mcp.identity import AgentIdentity
from app.core.orchestration.agent.recovery import RecoveryAgent


def identity():
    return AgentIdentity(
        subject="v54-recovery",
        roles=frozenset({"scheduler"}),
        business_domains=frozenset({"sales"}),
        environments=frozenset({"test"}),
        correlation_id="01900000-0000-7000-8000-000000000057",
    )


def schedule_plan(concurrency=4):
    return {
        "schema_version": "1.0",
        "timezone": "Asia/Shanghai",
        "triggers": [{"type": "manual"}],
        "max_concurrency": concurrency,
        "conflict_policy": "skip",
        "timeout_seconds": 600,
        "retry": {"max_attempts": 2, "delay_seconds": 30},
        "backfill": {"max_days": 1, "max_runs": 1},
    }


class Tools:
    def __init__(self):
        self.calls = []

    def call_tool(self, name, arguments):
        self.calls.append((name, copy.deepcopy(arguments)))
        return {
            "status": {
                "pause_schedule": "paused",
                "retry_failed_execution": "retry_started",
                "rollback_to_previous_version": "rolled_back",
            }[name]
        }


class Audit:
    def __init__(self):
        self.events = []

    def record(self, event):
        self.events.append(copy.deepcopy(event))


def test_pool_degradation_produces_deterministic_lower_concurrency_plan():
    tools = Tools()
    audit = Audit()
    agent = RecoveryAgent(scheduling_tools=tools, audit=audit)

    result = agent.recover(
        identity(),
        {
            "failure_kind": "datasource_pool_degraded",
            "business_domain": "sales",
            "environment": "test",
            "candidate_id": "candidate-v54",
            "schedule_plan": schedule_plan(4),
        },
    )

    assert result["status"] == "recovery_planned"
    assert result["action"] == "reduce_concurrency"
    assert result["schedule_plan"]["max_concurrency"] == 2
    assert result["schedule_plan"]["retry"]["max_attempts"] == 2
    assert tools.calls == []
    assert audit.events[-1]["decision"] == "allowed"


def test_transient_failure_uses_one_bounded_gateway_retry():
    tools = Tools()
    agent = RecoveryAgent(scheduling_tools=tools, audit=Audit(), max_retries=2)

    result = agent.recover(
        identity(),
        {
            "failure_kind": "transient_execution_failure",
            "business_domain": "sales",
            "environment": "test",
            "execution_id": "execution-v54",
            "retry_count": 1,
        },
    )

    assert result["status"] == "recovery_executed"
    assert result["action"] == "retry_failed_execution"
    assert tools.calls == [
        (
            "retry_failed_execution",
            {
                "execution_id": "execution-v54",
                "business_domain": "sales",
                "environment": "test",
                "max_retries": 2,
            },
        )
    ]


def test_exhausted_retry_rolls_back_to_previous_stable_version():
    tools = Tools()
    agent = RecoveryAgent(scheduling_tools=tools, audit=Audit(), max_retries=2)

    result = agent.recover(
        identity(),
        {
            "failure_kind": "transient_execution_failure",
            "business_domain": "sales",
            "environment": "test",
            "candidate_id": "candidate-v54",
            "execution_id": "execution-v54",
            "retry_count": 2,
        },
    )

    assert result["action"] == "rollback_to_previous_version"
    assert tools.calls[0][0] == "rollback_to_previous_version"
    assert tools.calls[0][1]["idempotency_key"].startswith("v54-recovery:")


def test_runner_or_datasource_outage_pauses_candidate():
    for failure_kind in ("runner_unavailable", "datasource_unavailable"):
        tools = Tools()
        agent = RecoveryAgent(scheduling_tools=tools, audit=Audit())

        result = agent.recover(
            identity(),
            {
                "failure_kind": failure_kind,
                "business_domain": "sales",
                "environment": "test",
                "candidate_id": "candidate-v54",
            },
        )

        assert result["action"] == "pause_schedule"
        assert tools.calls[0][0] == "pause_schedule"


def test_model_or_kestra_outage_keeps_published_fixed_schedule():
    for failure_kind in ("model_unavailable", "kestra_unavailable"):
        tools = Tools()
        agent = RecoveryAgent(scheduling_tools=tools, audit=Audit())

        result = agent.recover(
            identity(),
            {
                "failure_kind": failure_kind,
                "business_domain": "sales",
                "environment": "test",
            },
        )

        assert result["status"] == "degraded"
        assert result["action"] == "continue_published_schedule"
        assert tools.calls == []


def test_recovery_rejects_model_selected_or_dangerous_actions():
    agent = RecoveryAgent(scheduling_tools=Tools(), audit=Audit())

    with pytest.raises(ValueError, match="unsupported fields"):
        agent.recover(
            identity(),
            {
                "failure_kind": "transient_execution_failure",
                "business_domain": "sales",
                "environment": "test",
                "requested_action": "delete_execution",
            },
        )