Aether/tests/services/test_pool_manager_trace.py

"""Tests for pool manager trace data collection and attachment."""

from __future__ import annotations

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

import pytest

from src.services.provider.pool.config import PoolConfig
from src.services.provider.pool.manager import PoolManager
from src.services.provider.pool.trace import PoolSchedulingTrace


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


# ---------------------------------------------------------------------------
# Trace attachment to candidates
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_trace_attached_to_first_candidate() -> None:
    pool = PoolManager("prov-1", PoolConfig())
    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": 100.0, "key-2": 50.0},
        ),
    ):
        result = await pool.reorder_candidates("sess-1", [c1, c2])

    # First candidate should have _pool_scheduling_trace
    trace = getattr(result[0], "_pool_scheduling_trace", None)
    assert trace is not None
    assert isinstance(trace, PoolSchedulingTrace)
    assert trace.total_keys == 2
    assert trace.provider_id == "prov-1"


@pytest.mark.asyncio
async def test_pool_extra_data_on_selected_candidate() -> None:
    pool = PoolManager("prov-1", PoolConfig())
    c1 = _make_candidate("key-1")

    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},
        ),
        patch(
            "src.services.provider.pool.redis_ops.get_lru_scores",
            new_callable=AsyncMock,
            return_value={},
        ),
    ):
        result = await pool.reorder_candidates(None, [c1])

    extra = getattr(result[0], "_pool_extra_data", None)
    assert extra is not None
    assert "pool_selection" in extra


@pytest.mark.asyncio
async def test_pool_extra_data_on_skipped_candidate() -> None:
    pool = PoolManager("prov-1", PoolConfig())
    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 is skipped with pool_skip extra data
    skipped = [r for r in result if r.is_skipped]
    assert len(skipped) == 1
    extra = getattr(skipped[0], "_pool_extra_data", None)
    assert extra is not None
    assert extra["pool_skip"]["type"] == "cooldown"
    assert extra["pool_skip"]["cooldown_reason"] == "rate_limited_429"
    # TTL is not fetched on the scheduling hot path (include_ttl=False)
    assert "cooldown_ttl" not in extra["pool_skip"]


@pytest.mark.asyncio
async def test_trace_build_summary_matches() -> None:
    pool = PoolManager("prov-1", PoolConfig(cost_limit_per_key_tokens=1000))
    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": "overloaded_529",
                "key-2": None,
                "key-3": None,
            },
        ),
        patch(
            "src.services.provider.pool.redis_ops.batch_get_cost_totals",
            new_callable=AsyncMock,
            return_value={"key-2": 1500, "key-3": 200},
        ),
    ):
        result = await pool.reorder_candidates(None, [c1, c2, c3])

    trace = getattr(result[0], "_pool_scheduling_trace", None)
    assert trace is not None

    summary = trace.build_summary(success_key_id="key-3")
    assert summary["total_keys"] == 3
    assert summary["skipped_cooldown"] == 1  # key-1
    assert summary["skipped_cost"] == 1  # key-2
    assert summary["attempted"] == 1  # key-3
    assert summary["success_key_id"] == "key-3"[:8]


@pytest.mark.asyncio
async def test_trace_build_summary_uses_attempted_key_ids_when_provided() -> None:
    pool = PoolManager("prov-1", PoolConfig())
    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={},
        ),
    ):
        result = await pool.reorder_candidates(None, [c1, c2])

    trace = getattr(result[0], "_pool_scheduling_trace", None)
    assert trace is not None

    summary = trace.build_summary(
        success_key_id="key-1",
        attempted_key_ids={"key-1"},
    )
    assert summary["total_keys"] == 2
    assert summary["attempted"] == 1
    assert summary["success_key_id"] == "key-1"[:8]


@pytest.mark.asyncio
async def test_sticky_trace_info() -> None:
    pool = PoolManager("prov-1", PoolConfig(sticky_session_ttl_seconds=3600))
    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="key-2",
        ),
        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("sess-1", [c1, c2])

    trace = getattr(result[0], "_pool_scheduling_trace", None)
    assert trace is not None
    assert trace.sticky_session_used is True

    # key-2 (sticky hit) should be first
    assert result[0].key.id == "key-2"
    extra = getattr(result[0], "_pool_extra_data", None)
    assert extra is not None
    assert extra["pool_selection"]["sticky_hit"] is True