tests/test_cache.py

"""
Unit tests for CacheLayer.

Tests:
  - Exact cache hit/miss
  - Exact cache stores and retrieves correctly
  - Prefix cache hit on matching prefix
  - Prefix cache miss on different prefix
  - Dedup lock acquire / release
  - Dedup wait (subscriber) receives published result
"""
import asyncio
import json
import pytest
from unittest.mock import AsyncMock, MagicMock, patch


SAMPLE_REQ = {
    "model": "edge/auto",
    "messages": [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "Explain the concept of caching in distributed systems."},
    ]
}

SAMPLE_RESP = {
    "id": "test-completion-001",
    "object": "chat.completion",
    "model": "gpt-4o-mini",
    "choices": [{"index": 0, "message": {"role": "assistant", "content": "Caching stores frequently accessed data closer to the consumer..."}, "finish_reason": "stop"}],
    "usage": {"prompt_tokens": 25, "completion_tokens": 50, "total_tokens": 75, "estimated_cost_usd": 0.00003},
}

SAMPLE_RESP_2 = {
    "id": "test-completion-002",
    "object": "chat.completion",
    "model": "gpt-4o-mini",
    "choices": [{"index": 0, "message": {"role": "assistant", "content": "Different answer for a different request."}, "finish_reason": "stop"}],
    "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30, "estimated_cost_usd": 0.00001},
}


@pytest.fixture
def redis_store():
    """In-memory fake Redis for cache tests."""
    store: dict = {}
    pubsub_messages: dict = {}  # channel → list of messages

    redis = AsyncMock()

    async def get(key):
        return store.get(key)

    async def set(key, value, *args, **kwargs):
        store[key] = value

    async def delete(key):
        store.pop(key, None)

    async def zadd(key, mapping):
        if key not in store:
            store[key] = {}
        store[key].update(mapping)

    async def zremrangebyrank(key, start, end):
        pass

    async def expire(key, ttl):
        pass

    async def zrange(key, start, end):
        if key not in store or not isinstance(store[key], dict):
            return []
        return list(store[key].keys())

    async def set_nx(key, value, **kwargs):
        if key in store:
            return None
        store[key] = value
        return True

    redis.get.side_effect = get
    redis.set.side_effect = set
    redis.delete.side_effect = delete
    redis.zadd.side_effect = zadd
    redis.zremrangebyrank.side_effect = zremrangebyrank
    redis.expire.side_effect = expire
    redis.zrange.side_effect = zrange

    # nx=True set
    async def set_with_nx(key, value, nx=False, ex=None, **kwargs):
        if nx and key in store:
            return None
        store[key] = value
        return True

    redis.set.side_effect = set_with_nx

    return redis, store


@pytest.fixture
def cache(redis_store):
    from inferroute.cache import CacheLayer
    redis, store = redis_store
    layer = CacheLayer()
    with patch("inferroute.cache.get_redis_client", return_value=redis):
        yield layer, store, redis


@pytest.mark.asyncio
async def test_exact_cache_miss(cache):
    layer, store, redis = cache
    result = await layer.lookup_exact(SAMPLE_REQ)
    assert result is None


@pytest.mark.asyncio
async def test_exact_cache_store_and_hit(cache):
    layer, store, redis = cache
    await layer.store_exact(SAMPLE_REQ, SAMPLE_RESP)
    result = await layer.lookup_exact(SAMPLE_REQ)
    assert result is not None
    assert result["id"] == "test-completion-001"


@pytest.mark.asyncio
async def test_exact_cache_different_request_is_miss(cache):
    layer, store, redis = cache
    await layer.store_exact(SAMPLE_REQ, SAMPLE_RESP)

    different_req = {
        "model": "edge/auto",
        "messages": [{"role": "user", "content": "What is the weather today?"}]
    }
    result = await layer.lookup_exact(different_req)
    assert result is None


@pytest.mark.asyncio
async def test_exact_cache_ignores_routing_and_metadata_keys(cache):
    """Two requests that differ only in routing/metadata keys should share the cache."""
    layer, store, redis = cache
    req_with_meta = {**SAMPLE_REQ, "routing": {"policy": "cost"}, "metadata": {"user": "alice"}}

    await layer.store_exact(SAMPLE_REQ, SAMPLE_RESP)

    # Should still hit even with extra keys
    result = await layer.lookup_exact(req_with_meta)
    assert result is not None


@pytest.mark.asyncio
async def test_prefix_cache_disabled(cache):
    """Verify that lookup_prefix always returns None as prefix answer caching is disabled."""
    layer, store, redis = cache

    long_req = {
        "model": "edge/auto",
        "messages": [
            {"role": "user", "content": "Explain the concept of caching in distributed systems and its impact on scalability and performance in microservices architectures."}
        ]
    }
    await layer.store_exact(long_req, SAMPLE_RESP)

    # Request with a truncated prompt that still matches prefix
    prefix_req = {
        "model": "edge/auto",
        "messages": [
            {"role": "user", "content": "Explain the concept of caching in distributed systems"}
        ]
    }
    result = await layer.lookup_prefix(prefix_req)
    assert result is None


@pytest.mark.asyncio
async def test_dedup_lock_acquire_and_release(cache):
    layer, store, redis = cache
    is_owner = await layer.try_acquire_dedup_lock(SAMPLE_REQ)
    assert is_owner is True

    await layer.release_dedup_lock(SAMPLE_REQ)
    # After release, another caller can acquire
    is_owner_2 = await layer.try_acquire_dedup_lock(SAMPLE_REQ)
    assert is_owner_2 is True


@pytest.mark.asyncio
async def test_dedup_lock_second_caller_blocked(cache):
    """Second call with same request should see lock is taken."""
    layer, store, redis = cache
    is_owner_1 = await layer.try_acquire_dedup_lock(SAMPLE_REQ)
    assert is_owner_1 is True
    is_owner_2 = await layer.try_acquire_dedup_lock(SAMPLE_REQ)
    # Second caller should NOT be owner
    assert is_owner_2 is False


@pytest.mark.asyncio
async def test_cache_key_determinism(cache):
    """Same logical request should always produce the same cache key."""
    layer, _, _ = cache
    key1 = layer._exact_key(SAMPLE_REQ)
    key2 = layer._exact_key(dict(SAMPLE_REQ))
    assert key1 == key2
    assert key1.startswith("inferroute:cache:exact:")