Aether/tests/services/test_failover_engine.py

from __future__ import annotations

from types import SimpleNamespace
from typing import Any, AsyncIterator, cast
from unittest.mock import AsyncMock, MagicMock

import pytest

from src.models.database import Provider, ProviderAPIKey, ProviderEndpoint
from src.services.candidate.failover import FailoverEngine
from src.services.candidate.policy import RetryMode, RetryPolicy, SkipPolicy
from src.services.orchestration.error_classifier import ErrorAction, ErrorClassifier
from src.services.scheduling.schemas import PoolCandidate, ProviderCandidate
from src.services.task.core.protocol import AttemptKind, AttemptResult


def _make_candidate(
    *,
    provider_id: str = "p1",
    provider_name: str = "prov",
    endpoint_id: str = "e1",
    key_id: str = "k1",
    key_name: str = "key",
    auth_type: str = "api_key",
    priority: int = 0,
    is_cached: bool = False,
    is_skipped: bool = False,
    skip_reason: str | None = None,
    needs_conversion: bool = False,
    provider_max_retries: int | None = None,
    provider_config: dict[str, Any] | None = None,
) -> ProviderCandidate:
    provider = cast(
        Provider,
        SimpleNamespace(
            id=provider_id,
            name=provider_name,
            max_retries=provider_max_retries,
            config=provider_config,
        ),
    )
    endpoint = cast(ProviderEndpoint, SimpleNamespace(id=endpoint_id))
    key = cast(
        ProviderAPIKey,
        SimpleNamespace(id=key_id, name=key_name, auth_type=auth_type, priority=priority),
    )
    return ProviderCandidate(
        provider=provider,
        endpoint=endpoint,
        key=key,
        is_cached=is_cached,
        is_skipped=is_skipped,
        skip_reason=skip_reason,
        needs_conversion=needs_conversion,
    )


class _StubErrorClassifier:
    def __init__(self, *, action: ErrorAction, client_error: bool = False) -> None:
        self._action = action
        self._client_error = client_error

    def is_client_error(self, _text: str | None) -> bool:
        return self._client_error

    def classify(
        self, _error: Exception, *, has_retry_left: bool = False
    ) -> ErrorAction:  # noqa: ARG002
        return self._action


def _stub_classifier(*, action: ErrorAction, client_error: bool = False) -> ErrorClassifier:
    return cast(ErrorClassifier, _StubErrorClassifier(action=action, client_error=client_error))


@pytest.mark.asyncio
async def test_failover_engine_success_first_candidate() -> None:
    db = MagicMock()
    engine = FailoverEngine(db, error_classifier=_stub_classifier(action=ErrorAction.BREAK))

    candidates = [_make_candidate(provider_id="p1"), _make_candidate(provider_id="p2")]

    attempt = AsyncMock(
        return_value=AttemptResult(
            kind=AttemptKind.SYNC_RESPONSE,
            http_status=200,
            http_headers={},
            response_body={"ok": True},
        )
    )

    result = await engine.execute(
        candidates=candidates,
        attempt_func=attempt,
        retry_policy=RetryPolicy(mode=RetryMode.DISABLED, max_retries=1),
        skip_policy=SkipPolicy(),
        request_id=None,
    )

    assert result.success is True
    assert result.candidate_index == 0
    assert result.attempt_count == 1
    assert result.provider_id == "p1"
    assert result.response == {"ok": True}
    assert attempt.await_count == 1


@pytest.mark.asyncio
async def test_failover_engine_continue_to_next_candidate_on_error() -> None:
    db = MagicMock()
    engine = FailoverEngine(db, error_classifier=_stub_classifier(action=ErrorAction.BREAK))

    candidates = [_make_candidate(provider_id="p1"), _make_candidate(provider_id="p2")]

    attempt = AsyncMock(
        side_effect=[
            RuntimeError("boom"),
            AttemptResult(
                kind=AttemptKind.SYNC_RESPONSE,
                http_status=200,
                http_headers={},
                response_body={"ok": True},
            ),
        ]
    )

    result = await engine.execute(
        candidates=candidates,
        attempt_func=attempt,
        retry_policy=RetryPolicy(mode=RetryMode.DISABLED, max_retries=1),
        skip_policy=SkipPolicy(),
        request_id=None,
    )

    assert result.success is True
    assert result.candidate_index == 1
    assert result.provider_id == "p2"
    assert result.attempt_count == 2
    assert attempt.await_count == 2


@pytest.mark.asyncio
async def test_failover_engine_retry_same_candidate_when_classifier_says_continue() -> None:
    db = MagicMock()
    # ErrorAction.CONTINUE => retry current candidate (mapped to FailoverAction.RETRY)
    engine = FailoverEngine(db, error_classifier=_stub_classifier(action=ErrorAction.CONTINUE))

    candidates = [_make_candidate(provider_id="p1", is_cached=True, provider_max_retries=2)]

    attempt = AsyncMock(
        side_effect=[
            RuntimeError("transient"),
            AttemptResult(
                kind=AttemptKind.SYNC_RESPONSE,
                http_status=200,
                http_headers={},
                response_body={"ok": True},
            ),
        ]
    )

    result = await engine.execute(
        candidates=candidates,
        attempt_func=attempt,
        retry_policy=RetryPolicy(mode=RetryMode.ON_DEMAND, max_retries=2),
        skip_policy=SkipPolicy(),
        request_id=None,
    )

    assert result.success is True
    assert result.candidate_index == 0
    assert result.attempt_count == 2
    assert attempt.await_count == 2


@pytest.mark.asyncio
async def test_failover_engine_continues_when_classifier_raises() -> None:
    """After the 'default failover' change, RAISE no longer stops failover.
    All candidates should be attempted."""
    db = MagicMock()
    engine = FailoverEngine(db, error_classifier=_stub_classifier(action=ErrorAction.RAISE))

    candidates = [_make_candidate(provider_id="p1"), _make_candidate(provider_id="p2")]
    attempt = AsyncMock(side_effect=RuntimeError("client-ish"))

    result = await engine.execute(
        candidates=candidates,
        attempt_func=attempt,
        retry_policy=RetryPolicy(mode=RetryMode.DISABLED, max_retries=1),
        skip_policy=SkipPolicy(),
        request_id=None,
    )

    assert result.success is False
    assert result.error_type == "AllCandidatesFailed"
    # both candidates should be tried
    assert attempt.await_count == 2


async def _stream_two_chunks() -> AsyncIterator[bytes]:
    yield b"chunk1"
    yield b"chunk2"


async def _empty_stream() -> AsyncIterator[bytes]:
    if False:  # pragma: no cover
        yield b""
    return


@pytest.mark.asyncio
async def test_failover_engine_stream_probe_wraps_first_chunk() -> None:
    db = MagicMock()
    engine = FailoverEngine(db, error_classifier=_stub_classifier(action=ErrorAction.BREAK))

    candidates = [_make_candidate(provider_id="p1")]
    attempt = AsyncMock(
        return_value=AttemptResult(
            kind=AttemptKind.STREAM,
            http_status=200,
            http_headers={},
            stream_iterator=_stream_two_chunks(),
        )
    )

    result = await engine.execute(
        candidates=candidates,
        attempt_func=attempt,
        retry_policy=RetryPolicy(mode=RetryMode.DISABLED, max_retries=1),
        skip_policy=SkipPolicy(),
        request_id=None,
    )

    assert result.success is True
    assert result.attempt_result is not None
    assert result.attempt_result.kind == AttemptKind.STREAM

    collected: list[bytes] = []
    assert result.response is not None
    async for chunk in result.response:  # type: ignore[union-attr]
        collected.append(chunk)
    assert collected == [b"chunk1", b"chunk2"]


@pytest.mark.asyncio
async def test_failover_engine_stream_probe_empty_triggers_failover() -> None:
    db = MagicMock()
    engine = FailoverEngine(db, error_classifier=_stub_classifier(action=ErrorAction.BREAK))

    candidates = [_make_candidate(provider_id="p1"), _make_candidate(provider_id="p2")]

    attempt = AsyncMock(
        side_effect=[
            AttemptResult(
                kind=AttemptKind.STREAM,
                http_status=200,
                http_headers={},
                stream_iterator=_empty_stream(),
            ),
            AttemptResult(
                kind=AttemptKind.SYNC_RESPONSE,
                http_status=200,
                http_headers={},
                response_body={"ok": True},
            ),
        ]
    )

    result = await engine.execute(
        candidates=candidates,
        attempt_func=attempt,
        retry_policy=RetryPolicy(mode=RetryMode.DISABLED, max_retries=1),
        skip_policy=SkipPolicy(),
        request_id=None,
    )

    assert result.success is True
    assert result.candidate_index == 1
    assert result.attempt_count == 2


@pytest.mark.asyncio
async def test_failover_engine_pre_expand_marks_unused_slots_on_success(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    db = MagicMock()
    engine = FailoverEngine(db, error_classifier=_stub_classifier(action=ErrorAction.BREAK))

    # patch low-level record updater to observe unused marking
    engine._update_record = MagicMock()  # type: ignore[method-assign]
    engine.db.commit = MagicMock()  # type: ignore[method-assign]
    engine._commit_before_await = MagicMock()  # type: ignore[method-assign]

    c0 = _make_candidate(provider_id="p1", is_cached=True, provider_max_retries=2)
    c1 = _make_candidate(provider_id="p2", is_cached=False)

    attempt = AsyncMock(
        return_value=AttemptResult(
            kind=AttemptKind.SYNC_RESPONSE,
            http_status=200,
            http_headers={},
            response_body={"ok": True},
        )
    )

    record_map = {
        (0, 0): "r00",
        (0, 1): "r01",
        (1, 0): "r10",
    }

    result = await engine.execute(
        candidates=[c0, c1],
        attempt_func=attempt,
        retry_policy=RetryPolicy(mode=RetryMode.PRE_EXPAND, max_retries=2),
        skip_policy=SkipPolicy(),
        request_id=None,
        candidate_record_map=record_map,
    )

    assert result.success is True

    # Ensure we marked the remaining slots unused (r01 + r10)
    unused_record_ids = {
        call.args[0]
        for call in engine._update_record.call_args_list  # type: ignore[attr-defined]
        if call.kwargs.get("status") == "unused"
    }
    assert unused_record_ids == {"r01", "r10"}


# ========== error_stop_patterns with status_codes ==========


class _HttpError(Exception):
    """Stub exception with status_code and response text."""

    def __init__(self, status_code: int, text: str) -> None:
        super().__init__(text)
        self.status_code = status_code


@pytest.mark.asyncio
async def test_error_stop_pattern_with_matching_status_code_stops_failover() -> None:
    """When status_codes is set and matches, failover should stop."""
    db = MagicMock()
    engine = FailoverEngine(db, error_classifier=_stub_classifier(action=ErrorAction.BREAK))

    config = {
        "failover_rules": {
            "error_stop_patterns": [
                {"pattern": "content_policy", "status_codes": [403]},
            ],
        },
    }
    candidates = [
        _make_candidate(provider_id="p1", provider_config=config),
        _make_candidate(provider_id="p2"),
    ]

    attempt = AsyncMock(
        side_effect=[
            _HttpError(403, "content_policy_violation"),
            AttemptResult(
                kind=AttemptKind.SYNC_RESPONSE,
                http_status=200,
                http_headers={},
                response_body={"ok": True},
            ),
        ]
    )

    result = await engine.execute(
        candidates=candidates,
        attempt_func=attempt,
        retry_policy=RetryPolicy(mode=RetryMode.DISABLED, max_retries=1),
        skip_policy=SkipPolicy(),
        request_id=None,
    )

    # Should stop at first candidate, not try second
    assert result.success is False
    assert attempt.await_count == 1


@pytest.mark.asyncio
async def test_error_stop_pattern_with_non_matching_status_code_continues() -> None:
    """When status_codes is set but doesn't match, the rule is skipped and failover continues."""
    db = MagicMock()
    engine = FailoverEngine(db, error_classifier=_stub_classifier(action=ErrorAction.BREAK))

    config = {
        "failover_rules": {
            "error_stop_patterns": [
                {"pattern": "content_policy", "status_codes": [403]},
            ],
        },
    }
    candidates = [
        _make_candidate(provider_id="p1", provider_config=config),
        _make_candidate(provider_id="p2"),
    ]

    attempt = AsyncMock(
        side_effect=[
            _HttpError(500, "content_policy_violation"),
            AttemptResult(
                kind=AttemptKind.SYNC_RESPONSE,
                http_status=200,
                http_headers={},
                response_body={"ok": True},
            ),
        ]
    )

    result = await engine.execute(
        candidates=candidates,
        attempt_func=attempt,
        retry_policy=RetryPolicy(mode=RetryMode.DISABLED, max_retries=1),
        skip_policy=SkipPolicy(),
        request_id=None,
    )

    # status_code 500 doesn't match [403], so rule is skipped; failover continues to p2
    assert result.success is True
    assert result.provider_id == "p2"
    assert attempt.await_count == 2


@pytest.mark.asyncio
async def test_error_stop_pattern_without_status_codes_matches_any() -> None:
    """When status_codes is not set, the rule matches any status code (existing behavior)."""
    db = MagicMock()
    engine = FailoverEngine(db, error_classifier=_stub_classifier(action=ErrorAction.BREAK))

    config = {
        "failover_rules": {
            "error_stop_patterns": [
                {"pattern": "content_policy"},
            ],
        },
    }
    candidates = [
        _make_candidate(provider_id="p1", provider_config=config),
        _make_candidate(provider_id="p2"),
    ]

    attempt = AsyncMock(
        side_effect=[
            _HttpError(500, "content_policy_violation"),
            AttemptResult(
                kind=AttemptKind.SYNC_RESPONSE,
                http_status=200,
                http_headers={},
                response_body={"ok": True},
            ),
        ]
    )

    result = await engine.execute(
        candidates=candidates,
        attempt_func=attempt,
        retry_policy=RetryPolicy(mode=RetryMode.DISABLED, max_retries=1),
        skip_policy=SkipPolicy(),
        request_id=None,
    )

    # No status_codes filter, pattern matches -> stop
    assert result.success is False
    assert attempt.await_count == 1


@pytest.mark.asyncio
async def test_failover_engine_applies_backoff_every_tenth_failure(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    db = MagicMock()
    engine = FailoverEngine(db, error_classifier=_StubErrorClassifier(action=ErrorAction.BREAK))

    sleep_mock = AsyncMock()
    rotate_mock = AsyncMock(return_value=False)
    monkeypatch.setattr("src.services.candidate.failover.asyncio.sleep", sleep_mock)
    monkeypatch.setattr(engine, "_rotate_upstream_client", rotate_mock)

    await engine._apply_retry_pacing(
        candidate=_make_candidate(),
        consecutive_failures=10,
        error=RuntimeError("boom"),
        request_id="req-1",
    )

    sleep_mock.assert_awaited_once()
    rotate_mock.assert_not_awaited()


@pytest.mark.asyncio
async def test_failover_engine_rotates_client_on_stream_capacity_error(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    db = MagicMock()
    engine = FailoverEngine(db, error_classifier=_StubErrorClassifier(action=ErrorAction.BREAK))

    sleep_mock = AsyncMock()
    rotate_mock = AsyncMock(return_value=True)
    monkeypatch.setattr("src.services.candidate.failover.asyncio.sleep", sleep_mock)
    monkeypatch.setattr(engine, "_rotate_upstream_client", rotate_mock)

    await engine._apply_retry_pacing(
        candidate=_make_candidate(),
        consecutive_failures=3,
        error=RuntimeError("LocalProtocolError: Max outbound streams is 100, 100 open"),
        request_id="req-2",
    )

    rotate_mock.assert_awaited_once()
    sleep_mock.assert_awaited_once()


@pytest.mark.asyncio
async def test_failover_engine_stops_when_cancelled_before_attempt() -> None:
    db = MagicMock()
    engine = FailoverEngine(db, error_classifier=_StubErrorClassifier(action=ErrorAction.BREAK))

    candidates = [_make_candidate(provider_id="p1"), _make_candidate(provider_id="p2")]
    attempt = AsyncMock()

    async def _cancelled() -> bool:
        return True

    result = await engine.execute(
        candidates=candidates,
        attempt_func=attempt,
        retry_policy=RetryPolicy(mode=RetryMode.DISABLED, max_retries=1),
        skip_policy=SkipPolicy(),
        request_id=None,
        is_cancelled=_cancelled,
    )

    assert result.success is False
    assert result.error_type == "cancelled"
    assert result.last_status_code == 499
    assert result.attempt_count == 0
    assert attempt.await_count == 0
    assert [item.status for item in result.candidate_keys] == ["cancelled", "cancelled"]


@pytest.mark.asyncio
async def test_failover_engine_stops_before_next_pool_key_when_cancelled() -> None:
    db = MagicMock()
    engine = FailoverEngine(db, error_classifier=_StubErrorClassifier(action=ErrorAction.BREAK))

    provider = SimpleNamespace(id="p1", name="prov", max_retries=1, config={})
    endpoint = SimpleNamespace(id="e1")
    key1 = SimpleNamespace(id="k1", name="key-1", auth_type="api_key", priority=0)
    key2 = SimpleNamespace(id="k2", name="key-2", auth_type="api_key", priority=0)
    pool_candidate = PoolCandidate(
        provider=provider,  # type: ignore[arg-type]
        endpoint=endpoint,  # type: ignore[arg-type]
        key=key1,  # type: ignore[arg-type]
        pool_keys=[key1, key2],  # type: ignore[list-item]
    )

    attempt = AsyncMock(side_effect=RuntimeError("boom"))
    cancel_checks = {"count": 0}

    async def _cancelled() -> bool:
        cancel_checks["count"] += 1
        return cancel_checks["count"] >= 4

    result = await engine.execute(
        candidates=[pool_candidate],
        attempt_func=attempt,
        retry_policy=RetryPolicy(mode=RetryMode.DISABLED, max_retries=1),
        skip_policy=SkipPolicy(),
        request_id=None,
        is_cancelled=_cancelled,
    )

    assert result.success is False
    assert result.error_type == "cancelled"
    assert attempt.await_count == 1
    assert any(item.key_id == "k2" and item.status == "cancelled" for item in result.candidate_keys)