Aether/tests/services/test_pool_manager.py

"""Tests for pool_manager.py — candidate reordering, success/error hooks."""

from __future__ import annotations

from types import SimpleNamespace
from unittest.mock import AsyncMock, patch

import pytest

from src.services.provider.pool.config import PoolConfig, SchedulingPreset, ScoringWeights
from src.services.provider.pool.manager import PoolManager


def _make_candidate(
    key_id: str,
    *,
    is_skipped: bool = False,
    upstream_metadata: dict | None = None,
    oauth_invalid_reason: str | None = None,
) -> SimpleNamespace:
    return SimpleNamespace(
        key=SimpleNamespace(
            id=key_id,
            upstream_metadata=upstream_metadata,
            oauth_invalid_reason=oauth_invalid_reason,
        ),
        is_skipped=is_skipped,
        skip_reason=None,
    )


@pytest.fixture()
def pool() -> PoolManager:
    return PoolManager("provider-1", PoolConfig())


# ---------------------------------------------------------------------------
# reorder_candidates
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_reorder_empty_candidates(pool: PoolManager) -> None:
    result = await pool.reorder_candidates("sess-1", [])
    assert result == []


@pytest.mark.asyncio
async def test_reorder_sticky_hit_moves_to_front(pool: PoolManager) -> None:
    c1 = _make_candidate("key-1")
    c2 = _make_candidate("key-2")
    c3 = _make_candidate("key-3")

    with (
        patch(
            "src.services.provider.pool.redis_ops.get_sticky_binding",
            new_callable=AsyncMock,
            return_value="key-2",
        ),
        patch(
            "src.services.provider.pool.redis_ops.batch_get_cooldowns",
            new_callable=AsyncMock,
            return_value={"key-1": None, "key-2": None, "key-3": None},
        ),
        patch(
            "src.services.provider.pool.redis_ops.get_lru_scores",
            new_callable=AsyncMock,
            return_value={"key-1": 100.0, "key-2": 200.0, "key-3": 50.0},
        ),
    ):
        result = await pool.reorder_candidates("sess-1", [c1, c2, c3])

    assert result[0].key.id == "key-2"  # sticky hit first


@pytest.mark.asyncio
async def test_reorder_cooldown_keys_are_skipped(pool: PoolManager) -> None:
    c1 = _make_candidate("key-1")
    c2 = _make_candidate("key-2")

    with (
        patch(
            "src.services.provider.pool.redis_ops.get_sticky_binding",
            new_callable=AsyncMock,
            return_value=None,
        ),
        patch(
            "src.services.provider.pool.redis_ops.batch_get_cooldowns",
            new_callable=AsyncMock,
            return_value={"key-1": "rate_limited_429", "key-2": None},
        ),
        patch(
            "src.services.provider.pool.redis_ops.get_lru_scores",
            new_callable=AsyncMock,
            return_value={},
        ),
    ):
        result = await pool.reorder_candidates(None, [c1, c2])

    # key-1 should be skipped
    assert c1.is_skipped is True
    assert "cooldown" in (c1.skip_reason or "")
    # key-2 first in available
    assert result[0].key.id == "key-2"


@pytest.mark.asyncio
async def test_reorder_cost_exhausted_keys_are_skipped() -> None:
    pool = PoolManager(
        "provider-1",
        PoolConfig(cost_limit_per_key_tokens=1000, lru_enabled=False),
    )
    c1 = _make_candidate("key-1")
    c2 = _make_candidate("key-2")

    with (
        patch(
            "src.services.provider.pool.redis_ops.get_sticky_binding",
            new_callable=AsyncMock,
            return_value=None,
        ),
        patch(
            "src.services.provider.pool.redis_ops.batch_get_cooldowns",
            new_callable=AsyncMock,
            return_value={"key-1": None, "key-2": None},
        ),
        patch(
            "src.services.provider.pool.redis_ops.batch_get_cost_totals",
            new_callable=AsyncMock,
            return_value={"key-1": 1500, "key-2": 200},
        ),
    ):
        result = await pool.reorder_candidates(None, [c1, c2])

    assert c1.is_skipped is True
    assert "cost" in (c1.skip_reason or "")
    assert result[0].key.id == "key-2"


@pytest.mark.asyncio
async def test_reorder_account_blocked_keys_are_skipped() -> None:
    pool = PoolManager("provider-1", PoolConfig(), provider_type="kiro")
    c1 = _make_candidate(
        "key-1",
        upstream_metadata={"kiro": {"is_banned": True, "ban_reason": "account suspended"}},
    )
    c2 = _make_candidate("key-2")

    with (
        patch(
            "src.services.provider.pool.redis_ops.get_sticky_binding",
            new_callable=AsyncMock,
            return_value=None,
        ),
        patch(
            "src.services.provider.pool.redis_ops.batch_get_cooldowns",
            new_callable=AsyncMock,
            return_value={"key-1": None, "key-2": None},
        ),
        patch(
            "src.services.provider.pool.redis_ops.get_lru_scores",
            new_callable=AsyncMock,
            return_value={},
        ),
    ):
        result = await pool.reorder_candidates(None, [c1, c2])

    assert c1.is_skipped is True
    assert "account blocked" in (c1.skip_reason or "")
    assert result[0].key.id == "key-2"
    assert c1._pool_extra_data["pool_skip"]["type"] == "account_blocked"
    assert c1._pool_extra_data["pool_skip"]["account_block_label"] == "账号封禁"


@pytest.mark.asyncio
async def test_reorder_multi_score_uses_composite_score() -> None:
    pool = PoolManager(
        "provider-1",
        PoolConfig(
            scheduling_mode="multi_score",
            strategies=("multi_score",),
            scoring_weights=ScoringWeights(lru=0.0, latency=1.0, health=0.0, cost_remaining=0.0),
        ),
    )
    c1 = _make_candidate("key-1")
    c2 = _make_candidate("key-2")

    with (
        patch(
            "src.services.provider.pool.redis_ops.get_sticky_binding",
            new_callable=AsyncMock,
            return_value=None,
        ),
        patch(
            "src.services.provider.pool.redis_ops.batch_get_cooldowns",
            new_callable=AsyncMock,
            return_value={"key-1": None, "key-2": None},
        ),
        patch(
            "src.services.provider.pool.redis_ops.get_lru_scores",
            new_callable=AsyncMock,
            return_value={"key-1": 10.0, "key-2": 20.0},
        ),
        patch(
            "src.services.provider.pool.redis_ops.batch_get_latency_avgs",
            new_callable=AsyncMock,
            return_value={"key-1": 600.0, "key-2": 120.0},
        ),
    ):
        result = await pool.reorder_candidates(None, [c1, c2])

    assert result[0].key.id == "key-2"
    assert result[0]._pool_extra_data["pool_selection"]["reason"] == "multi_score"
    assert result[0]._pool_extra_data["pool_selection"]["scoring_mode"] == "multi_score"


@pytest.mark.asyncio
async def test_reorder_multi_score_presets_free_team_first() -> None:
    pool = PoolManager(
        "provider-1",
        PoolConfig(
            scheduling_mode="multi_score",
            strategies=("multi_score",),
            scheduling_presets=(
                SchedulingPreset(preset="free_team_first", enabled=True, mode="both"),
            ),
        ),
    )
    c1 = _make_candidate("key-1", upstream_metadata={"codex": {"plan_type": "plus"}})
    c2 = _make_candidate("key-2", upstream_metadata={"codex": {"plan_type": "team"}})

    with (
        patch(
            "src.services.provider.pool.redis_ops.get_sticky_binding",
            new_callable=AsyncMock,
            return_value=None,
        ),
        patch(
            "src.services.provider.pool.redis_ops.batch_get_cooldowns",
            new_callable=AsyncMock,
            return_value={"key-1": None, "key-2": None},
        ),
        patch(
            "src.services.provider.pool.redis_ops.get_lru_scores",
            new_callable=AsyncMock,
            return_value={"key-1": 100.0, "key-2": 100.0},
        ),
        patch(
            "src.services.provider.pool.redis_ops.batch_get_latency_avgs",
            new_callable=AsyncMock,
            return_value={},
        ),
    ):
        result = await pool.reorder_candidates(None, [c1, c2])

    assert result[0].key.id == "key-2"


@pytest.mark.asyncio
async def test_reorder_lru_sorts_least_recently_used_first(
    pool: PoolManager,
) -> None:
    c1 = _make_candidate("key-1")
    c2 = _make_candidate("key-2")
    c3 = _make_candidate("key-3")

    with (
        patch(
            "src.services.provider.pool.redis_ops.get_sticky_binding",
            new_callable=AsyncMock,
            return_value=None,
        ),
        patch(
            "src.services.provider.pool.redis_ops.batch_get_cooldowns",
            new_callable=AsyncMock,
            return_value={"key-1": None, "key-2": None, "key-3": None},
        ),
        patch(
            "src.services.provider.pool.redis_ops.get_lru_scores",
            new_callable=AsyncMock,
            return_value={"key-1": 300.0, "key-2": 100.0, "key-3": 200.0},
        ),
    ):
        result = await pool.reorder_candidates(None, [c1, c2, c3])

    # Least recently used (lowest score) first
    available = [r for r in result if not r.is_skipped]
    assert available[0].key.id == "key-2"
    assert available[1].key.id == "key-3"
    assert available[2].key.id == "key-1"


# ---------------------------------------------------------------------------
# on_request_success
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_on_request_success_binds_sticky_and_touches_lru(
    pool: PoolManager,
) -> None:
    with (
        patch(
            "src.services.provider.pool.redis_ops.set_sticky_binding",
            new_callable=AsyncMock,
        ) as mock_sticky,
        patch(
            "src.services.provider.pool.redis_ops.touch_lru",
            new_callable=AsyncMock,
        ) as mock_lru,
        patch(
            "src.services.provider.pool.redis_ops.add_cost_entry",
            new_callable=AsyncMock,
        ) as mock_cost,
    ):
        await pool.on_request_success(session_uuid="sess-1", key_id="key-1", tokens_used=0)

    mock_sticky.assert_called_once_with("provider-1", "sess-1", "key-1", 3600)
    mock_lru.assert_called_once_with("provider-1", "key-1")
    mock_cost.assert_not_called()  # tokens_used=0


@pytest.mark.asyncio
async def test_on_request_success_records_cost_when_configured() -> None:
    pool = PoolManager(
        "provider-1",
        PoolConfig(cost_limit_per_key_tokens=50000),
    )
    with (
        patch(
            "src.services.provider.pool.redis_ops.set_sticky_binding",
            new_callable=AsyncMock,
        ),
        patch(
            "src.services.provider.pool.redis_ops.touch_lru",
            new_callable=AsyncMock,
        ),
        patch(
            "src.services.provider.pool.redis_ops.add_cost_entry",
            new_callable=AsyncMock,
        ) as mock_cost,
    ):
        await pool.on_request_success(session_uuid="sess-1", key_id="key-1", tokens_used=500)

    mock_cost.assert_called_once_with("provider-1", "key-1", 500, 18000)


@pytest.mark.asyncio
async def test_on_request_success_records_latency_when_multi_score_enabled() -> None:
    pool = PoolManager(
        "provider-1",
        PoolConfig(
            scheduling_mode="multi_score",
            latency_window_seconds=7200,
            latency_sample_limit=80,
        ),
    )

    with (
        patch(
            "src.services.provider.pool.redis_ops.set_sticky_binding",
            new_callable=AsyncMock,
        ),
        patch(
            "src.services.provider.pool.redis_ops.touch_lru",
            new_callable=AsyncMock,
        ),
        patch(
            "src.services.provider.pool.redis_ops.record_latency",
            new_callable=AsyncMock,
        ) as mock_latency,
    ):
        await pool.on_request_success(
            session_uuid=None,
            key_id="key-1",
            tokens_used=0,
            ttfb_ms=321,
        )

    mock_latency.assert_called_once_with("provider-1", "key-1", 321, 7200, 80)


# ---------------------------------------------------------------------------
# on_request_error
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_on_request_error_delegates_to_health_policy(
    pool: PoolManager,
) -> None:
    with patch(
        "src.services.provider.pool.health_policy.apply_health_policy",
        new_callable=AsyncMock,
    ) as mock_hp:
        await pool.on_request_error(
            key_id="key-1", status_code=429, error_body=None, response_headers=None
        )

    mock_hp.assert_called_once()
    call_kwargs = mock_hp.call_args.kwargs
    assert call_kwargs["key_id"] == "key-1"
    assert call_kwargs["status_code"] == 429


# ---------------------------------------------------------------------------
# is_key_schedulable
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_is_key_schedulable_returns_false_when_in_cooldown(
    pool: PoolManager,
) -> None:
    with patch(
        "src.services.provider.pool.redis_ops.get_cooldown",
        new_callable=AsyncMock,
        return_value="rate_limited_429",
    ):
        ok, reason = await pool.is_key_schedulable("key-1")

    assert ok is False
    assert "cooldown" in (reason or "")


@pytest.mark.asyncio
async def test_is_key_schedulable_returns_true_when_healthy(
    pool: PoolManager,
) -> None:
    with patch(
        "src.services.provider.pool.redis_ops.get_cooldown",
        new_callable=AsyncMock,
        return_value=None,
    ):
        ok, reason = await pool.is_key_schedulable("key-1")

    assert ok is True
    assert reason is None