Aether/tests/services/test_cache_affinity_manager.py

import json
from fnmatch import fnmatch

import pytest

from src.services.scheduling.affinity_manager import CacheAffinityManager


class _FakeRedis:
    def __init__(self, store: dict[str, str]):
        self.store = dict(store)
        self.keys_called = 0
        self.scan_calls = 0
        self.unlink_batches: list[tuple[str, ...]] = []
        self.delete_batches: list[tuple[str, ...]] = []

    async def scan(  # type: ignore[override]
        self, cursor: int = 0, match: str | None = None, count: int = 10
    ) -> tuple[int, list[str]]:
        self.scan_calls += 1
        pattern = match or "*"
        matched = [key for key in sorted(self.store) if fnmatch(key, pattern)]
        start = int(cursor)
        batch = matched[start : start + count]
        next_cursor = 0 if start + count >= len(matched) else start + count
        return next_cursor, batch

    async def get(self, key: str) -> str | None:
        return self.store.get(key)

    async def mget(self, keys: list[str]) -> list[str | None]:
        return [self.store.get(k) for k in keys]

    async def unlink(self, *keys: str) -> int:
        self.unlink_batches.append(tuple(keys))
        deleted = 0
        for key in keys:
            if key in self.store:
                deleted += 1
                self.store.pop(key, None)
        return deleted

    async def delete(self, *keys: str) -> int:
        self.delete_batches.append(tuple(keys))
        deleted = 0
        for key in keys:
            if key in self.store:
                deleted += 1
                self.store.pop(key, None)
        return deleted

    async def keys(self, _pattern: str) -> list[str]:
        self.keys_called += 1
        raise AssertionError("clear operations should not call Redis KEYS")


@pytest.mark.asyncio
async def test_clear_all_uses_scan_and_clears_l1_cache() -> None:
    redis = _FakeRedis(
        {
            "cache_affinity:user-1:openai:model-a": json.dumps({"provider_id": "provider-a"}),
            "cache_affinity:user-2:openai:model-b": json.dumps({"provider_id": "provider-b"}),
            "cache_affinity:user-3:claude:model-c": json.dumps({"provider_id": "provider-c"}),
            "other:key": json.dumps({"provider_id": "provider-x"}),
        }
    )
    manager = CacheAffinityManager(redis_client=redis)

    await manager._set_l1_entry(
        "cache_affinity:user-1:openai:model-a", {"provider_id": "provider-a"}
    )
    await manager._set_l1_entry("other:key", {"provider_id": "provider-x"})

    deleted_count = await manager.clear_all()

    assert deleted_count == 3
    assert redis.keys_called == 0
    assert redis.scan_calls >= 1
    assert redis.unlink_batches == [
        (
            "cache_affinity:user-1:openai:model-a",
            "cache_affinity:user-2:openai:model-b",
            "cache_affinity:user-3:claude:model-c",
        )
    ]
    assert "other:key" in redis.store
    assert all(not key.startswith("cache_affinity:") for key in redis.store)
    assert await manager._get_l1_entry("cache_affinity:user-1:openai:model-a") is None
    assert await manager._get_l1_entry("other:key") == {"provider_id": "provider-x"}


@pytest.mark.asyncio
async def test_invalidate_all_for_provider_uses_scan_not_keys() -> None:
    redis = _FakeRedis(
        {
            "cache_affinity:user-1:openai:model-a": json.dumps({"provider_id": "provider-a"}),
            "cache_affinity:user-2:openai:model-b": json.dumps({"provider_id": "provider-b"}),
            "cache_affinity:user-3:claude:model-c": json.dumps({"provider_id": "provider-a"}),
        }
    )
    manager = CacheAffinityManager(redis_client=redis)

    await manager._set_l1_entry(
        "cache_affinity:user-1:openai:model-a", {"provider_id": "provider-a"}
    )
    await manager._set_l1_entry(
        "cache_affinity:user-2:openai:model-b", {"provider_id": "provider-b"}
    )
    await manager._set_l1_entry(
        "cache_affinity:user-3:claude:model-c", {"provider_id": "provider-a"}
    )

    deleted_count = await manager.invalidate_all_for_provider("provider-a")

    assert deleted_count == 2
    assert redis.keys_called == 0
    assert redis.scan_calls >= 1
    # Pipeline path: batch MGET + batch UNLINK (not per-key DELETE).
    assert len(redis.unlink_batches) == 1
    assert set(redis.unlink_batches[0]) == {
        "cache_affinity:user-1:openai:model-a",
        "cache_affinity:user-3:claude:model-c",
    }
    assert redis.delete_batches == []
    assert "cache_affinity:user-2:openai:model-b" in redis.store
    assert "cache_affinity:user-1:openai:model-a" not in redis.store
    assert "cache_affinity:user-3:claude:model-c" not in redis.store
    assert await manager._get_l1_entry("cache_affinity:user-1:openai:model-a") is None
    assert await manager._get_l1_entry("cache_affinity:user-2:openai:model-b") == {
        "provider_id": "provider-b"
    }