Aether/tests/services/test_pool_health_policy.py

"""Tests for pool_health_policy.py — error code classification."""

from __future__ import annotations

import json
from unittest.mock import AsyncMock, patch

import pytest

from src.services.provider.pool.config import (
    PoolConfig,
    UnschedulableRule,
)
from src.services.provider.pool.health_policy import (
    _extract_error_message,
    _parse_retry_after,
    apply_health_policy,
)

PID = "provider-test"
KID = "key-test"


@pytest.fixture()
def config() -> PoolConfig:
    return PoolConfig()


# ---------------------------------------------------------------------------
# Helper tests
# ---------------------------------------------------------------------------


def test_parse_retry_after_none() -> None:
    assert _parse_retry_after(None) is None
    assert _parse_retry_after({}) is None


def test_parse_retry_after_valid() -> None:
    assert _parse_retry_after({"retry-after": "60"}) == 60
    assert _parse_retry_after({"Retry-After": "120"}) == 120


def test_parse_retry_after_clamped() -> None:
    assert _parse_retry_after({"retry-after": "0"}) == 1
    assert _parse_retry_after({"retry-after": "9999"}) == 3600


def test_parse_retry_after_invalid() -> None:
    assert _parse_retry_after({"retry-after": "not-a-number"}) is None


def test_extract_error_message_json_error_object() -> None:
    body = json.dumps({"error": {"message": "bad request"}})
    assert _extract_error_message(body) == "bad request"


def test_extract_error_message_json_error_string() -> None:
    body = json.dumps({"error": "something went wrong"})
    assert _extract_error_message(body) == "something went wrong"


def test_extract_error_message_plain_text() -> None:
    assert _extract_error_message("plain text error") == "plain text error"


def test_extract_error_message_empty() -> None:
    assert _extract_error_message(None) == ""
    assert _extract_error_message("") == ""


# ---------------------------------------------------------------------------
# Status code handling
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_401_invalidates_oauth_cache_and_sets_cooldown(config: PoolConfig) -> None:
    with (
        patch(
            "src.services.provider.pool.redis_ops.invalidate_oauth_token_cache",
            new_callable=AsyncMock,
        ) as mock_inv,
        patch(
            "src.services.provider.pool.redis_ops.set_cooldown",
            new_callable=AsyncMock,
        ) as mock_cd,
    ):
        await apply_health_policy(
            provider_id=PID,
            key_id=KID,
            status_code=401,
            error_body=None,
            response_headers=None,
            config=config,
        )

    mock_inv.assert_called_once_with(KID)
    mock_cd.assert_called_once_with(PID, KID, "auth_failed_401", ttl=60)


@pytest.mark.asyncio
async def test_402_sets_long_cooldown(config: PoolConfig) -> None:
    with patch(
        "src.services.provider.pool.redis_ops.set_cooldown",
        new_callable=AsyncMock,
    ) as mock_cd:
        await apply_health_policy(
            provider_id=PID,
            key_id=KID,
            status_code=402,
            error_body=None,
            response_headers=None,
            config=config,
        )

    mock_cd.assert_called_once_with(PID, KID, "payment_required_402", ttl=3600)


@pytest.mark.asyncio
async def test_403_default_sets_medium_cooldown(config: PoolConfig) -> None:
    with patch(
        "src.services.provider.pool.redis_ops.set_cooldown",
        new_callable=AsyncMock,
    ) as mock_cd:
        await apply_health_policy(
            provider_id=PID,
            key_id=KID,
            status_code=403,
            error_body=None,
            response_headers=None,
            config=config,
        )

    mock_cd.assert_called_once_with(PID, KID, "forbidden_403", ttl=300)


@pytest.mark.asyncio
async def test_403_suspended_body_sets_long_cooldown(config: PoolConfig) -> None:
    with patch(
        "src.services.provider.pool.redis_ops.set_cooldown",
        new_callable=AsyncMock,
    ) as mock_cd:
        await apply_health_policy(
            provider_id=PID,
            key_id=KID,
            status_code=403,
            error_body=json.dumps({"error": {"message": "account suspended"}}),
            response_headers=None,
            config=config,
        )

    mock_cd.assert_called_once_with(PID, KID, "forbidden_403", ttl=3600)


@pytest.mark.asyncio
async def test_400_with_org_disabled_pattern(config: PoolConfig) -> None:
    body = json.dumps({"error": {"message": "Your organization has been disabled"}})
    with patch(
        "src.services.provider.pool.redis_ops.set_cooldown",
        new_callable=AsyncMock,
    ) as mock_cd:
        await apply_health_policy(
            provider_id=PID,
            key_id=KID,
            status_code=400,
            error_body=body,
            response_headers=None,
            config=config,
        )

    mock_cd.assert_called_once()
    assert "account_disabled_400" in mock_cd.call_args.args[2]


@pytest.mark.asyncio
async def test_400_without_pattern_does_nothing(config: PoolConfig) -> None:
    body = json.dumps({"error": {"message": "invalid json field"}})
    with patch(
        "src.services.provider.pool.redis_ops.set_cooldown",
        new_callable=AsyncMock,
    ) as mock_cd:
        await apply_health_policy(
            provider_id=PID,
            key_id=KID,
            status_code=400,
            error_body=body,
            response_headers=None,
            config=config,
        )

    mock_cd.assert_not_called()


@pytest.mark.asyncio
async def test_429_uses_retry_after_header(config: PoolConfig) -> None:
    with patch(
        "src.services.provider.pool.redis_ops.set_cooldown",
        new_callable=AsyncMock,
    ) as mock_cd:
        await apply_health_policy(
            provider_id=PID,
            key_id=KID,
            status_code=429,
            error_body=None,
            response_headers={"retry-after": "120"},
            config=config,
        )

    mock_cd.assert_called_once_with(PID, KID, "rate_limited_429", ttl=120)


@pytest.mark.asyncio
async def test_429_falls_back_to_config_default(config: PoolConfig) -> None:
    with patch(
        "src.services.provider.pool.redis_ops.set_cooldown",
        new_callable=AsyncMock,
    ) as mock_cd:
        await apply_health_policy(
            provider_id=PID,
            key_id=KID,
            status_code=429,
            error_body=None,
            response_headers=None,
            config=config,
        )

    mock_cd.assert_called_once_with(PID, KID, "rate_limited_429", ttl=300)


@pytest.mark.asyncio
async def test_529_uses_overload_cooldown(config: PoolConfig) -> None:
    with patch(
        "src.services.provider.pool.redis_ops.set_cooldown",
        new_callable=AsyncMock,
    ) as mock_cd:
        await apply_health_policy(
            provider_id=PID,
            key_id=KID,
            status_code=529,
            error_body=None,
            response_headers=None,
            config=config,
        )

    mock_cd.assert_called_once_with(PID, KID, "overloaded_529", ttl=30)


@pytest.mark.asyncio
async def test_503_uses_retry_after_when_present(config: PoolConfig) -> None:
    with patch(
        "src.services.provider.pool.redis_ops.set_cooldown",
        new_callable=AsyncMock,
    ) as mock_cd:
        await apply_health_policy(
            provider_id=PID,
            key_id=KID,
            status_code=503,
            error_body=None,
            response_headers={"retry-after": "45"},
            config=config,
        )

    mock_cd.assert_called_once_with(PID, KID, "service_unavailable_503", ttl=45)


@pytest.mark.asyncio
async def test_500_uses_overload_cooldown(config: PoolConfig) -> None:
    with patch(
        "src.services.provider.pool.redis_ops.set_cooldown",
        new_callable=AsyncMock,
    ) as mock_cd:
        await apply_health_policy(
            provider_id=PID,
            key_id=KID,
            status_code=500,
            error_body=None,
            response_headers=None,
            config=config,
        )

    mock_cd.assert_called_once_with(PID, KID, "server_error_500", ttl=30)


# ---------------------------------------------------------------------------
# Unschedulable keyword rules
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_keyword_rule_matches_and_sets_cooldown() -> None:
    cfg = PoolConfig(
        unschedulable_rules=[UnschedulableRule(keyword="capacity", duration_minutes=10)]
    )
    body = json.dumps({"error": {"message": "Server at capacity, try later"}})
    with patch(
        "src.services.provider.pool.redis_ops.set_cooldown",
        new_callable=AsyncMock,
    ) as mock_cd:
        await apply_health_policy(
            provider_id=PID,
            key_id=KID,
            status_code=500,
            error_body=body,
            response_headers=None,
            config=cfg,
        )

    mock_cd.assert_called_once_with(PID, KID, "rule:capacity", ttl=600)


# ---------------------------------------------------------------------------
# health_policy_enabled = False
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_disabled_health_policy_does_nothing() -> None:
    cfg = PoolConfig(health_policy_enabled=False)
    with patch(
        "src.services.provider.pool.redis_ops.set_cooldown",
        new_callable=AsyncMock,
    ) as mock_cd:
        await apply_health_policy(
            provider_id=PID,
            key_id=KID,
            status_code=429,
            error_body=None,
            response_headers=None,
            config=cfg,
        )

    mock_cd.assert_not_called()