litellm/tests/unit/caching/test_in_memory_cache.py

import asyncio
import json
import random
import sys
import threading
import time
from collections.abc import Iterator
from concurrent.futures import ThreadPoolExecutor
from typing import Final
from unittest.mock import MagicMock, patch

import httpx
import pytest
import respx
from fastapi.testclient import TestClient

from unittest.mock import AsyncMock

from litellm.caching.in_memory_cache import InMemoryCache


class _SlowInt(int):
    def __add__(self, value: int) -> "_SlowInt":
        time.sleep(0.05)
        return _SlowInt(int(self) + value)


def test_increment_cache_is_atomic_under_thread_concurrency():
    cache = InMemoryCache()
    seed = 1000
    cache.set_cache("counter", _SlowInt(seed))
    thread_count = 8
    barrier = threading.Barrier(thread_count)

    def increment(_: int) -> float:
        barrier.wait()
        return cache.increment_cache("counter", 1)

    with ThreadPoolExecutor(max_workers=thread_count) as executor:
        tuple(executor.map(increment, range(thread_count)))

    assert cache.get_cache("counter") == seed + thread_count


async def test_async_increment_delegates_to_locked_sync_path():
    cache = InMemoryCache()
    assert await cache.async_increment("counter", 2) == 2
    assert await cache.async_increment("counter", 3) == 5
    assert cache.get_cache("counter") == 5


def test_in_memory_openai_obj_cache():
    from openai import OpenAI

    openai_obj = OpenAI(api_key="my-fake-key")

    in_memory_cache = InMemoryCache()

    in_memory_cache.set_cache(key="my-fake-key", value=openai_obj)

    cached_obj = in_memory_cache.get_cache(key="my-fake-key")

    assert cached_obj is not None

    assert cached_obj == openai_obj


def test_in_memory_cache_max_size_per_item():
    """
    Test that the cache will not store items larger than the max size per item
    """
    in_memory_cache = InMemoryCache(max_size_per_item=100)

    result = in_memory_cache.check_value_size("a" * 100000000)

    assert result is False


def test_in_memory_cache_ttl():
    """
    Check that
    - if ttl is not set, it will be set to default ttl
    - if object expires, the ttl is also removed
    """
    in_memory_cache = InMemoryCache()

    in_memory_cache.set_cache(key="my-fake-key", value="my-fake-value", ttl=10)
    initial_ttl_time = in_memory_cache.ttl_dict["my-fake-key"]
    assert initial_ttl_time is not None

    in_memory_cache.set_cache(key="my-fake-key", value="my-fake-value-2", ttl=10)
    new_ttl_time = in_memory_cache.ttl_dict["my-fake-key"]
    assert new_ttl_time == initial_ttl_time  # ttl should not be updated

    ## On object expiration, the ttl should be removed
    in_memory_cache.set_cache(key="new-fake-key", value="new-fake-value", ttl=1)
    new_ttl_time = in_memory_cache.ttl_dict["new-fake-key"]
    assert new_ttl_time is not None
    time.sleep(1)
    cached_obj = in_memory_cache.get_cache(key="new-fake-key")
    new_ttl_time = in_memory_cache.ttl_dict.get("new-fake-key")
    assert new_ttl_time is None


def test_in_memory_cache_ttl_allow_override():
    """
    Check that
    - if ttl is not set, it will be set to default ttl
    - if object expires, the ttl is also removed
    """
    in_memory_cache = InMemoryCache()
    ## On object expiration, but no get_cache, the override should be allowed
    in_memory_cache.set_cache(key="new-fake-key", value="new-fake-value", ttl=1)
    initial_ttl_time = in_memory_cache.ttl_dict["new-fake-key"]
    assert initial_ttl_time is not None
    time.sleep(1)

    in_memory_cache.set_cache(key="new-fake-key", value="new-fake-value-2", ttl=1)
    new_ttl_time = in_memory_cache.ttl_dict["new-fake-key"]
    assert new_ttl_time is not None
    assert new_ttl_time != initial_ttl_time


def test_in_memory_cache_max_size_with_ttl():
    """
    Test that max_size_in_memory is respected even when all items have long TTLs.
    This tests the fix for the unbounded growth issue.
    """
    in_memory_cache = InMemoryCache(max_size_in_memory=3)
    long_ttl = 86400  # 1 day

    # Fill the cache to max capacity
    for i in range(3):
        in_memory_cache.set_cache(key=f"key_{i}", value=f"value_{i}", ttl=long_ttl)
        time.sleep(0.01)  # Small delay to ensure different timestamps

    assert len(in_memory_cache.cache_dict) == 3
    assert len(in_memory_cache.ttl_dict) == 3

    # Add another item - should evict the earliest item
    in_memory_cache.set_cache(key="key_3", value="value_3", ttl=long_ttl)

    # Cache should still be at max size, not larger
    assert len(in_memory_cache.cache_dict) == 3
    assert len(in_memory_cache.ttl_dict) == 3

    # key_0 should have been evicted (it was added first)
    assert "key_0" not in in_memory_cache.cache_dict
    assert "key_0" not in in_memory_cache.ttl_dict

    # Other keys should still be present
    assert "key_1" in in_memory_cache.cache_dict
    assert "key_2" in in_memory_cache.cache_dict
    assert "key_3" in in_memory_cache.cache_dict


def test_in_memory_cache_expired_items_evicted_first():
    """
    Test that expired items are evicted before non-expired items when cache is full.
    """
    in_memory_cache = InMemoryCache(max_size_in_memory=3)

    # Add items with short TTL that will expire
    in_memory_cache.set_cache(key="expired_1", value="value_1", ttl=1)
    in_memory_cache.set_cache(key="expired_2", value="value_2", ttl=1)

    # Add item with long TTL
    in_memory_cache.set_cache(key="long_lived", value="value_long", ttl=86400)

    assert len(in_memory_cache.cache_dict) == 3

    # Wait for short TTL items to expire
    time.sleep(2)

    # Add new item - should evict expired items first, not the long-lived one
    in_memory_cache.set_cache(key="new_item", value="new_value", ttl=86400)

    # Long-lived item should still be present
    assert "long_lived" in in_memory_cache.cache_dict
    assert "new_item" in in_memory_cache.cache_dict

    # Expired items should be gone
    assert "expired_1" not in in_memory_cache.cache_dict
    assert "expired_2" not in in_memory_cache.cache_dict
    assert "expired_1" not in in_memory_cache.ttl_dict
    assert "expired_2" not in in_memory_cache.ttl_dict


def test_in_memory_cache_eviction_order():
    """
    Test that when non-expired items need to be evicted, those with earliest expiration times are evicted first.
    """
    in_memory_cache = InMemoryCache(max_size_in_memory=2)

    # Add items with different TTLs
    now = time.time()
    in_memory_cache.set_cache(
        key="early_expire", value="value_1", ttl=100
    )  # expires in 100 seconds
    time.sleep(0.01)
    in_memory_cache.set_cache(
        key="late_expire", value="value_2", ttl=200
    )  # expires in 200 seconds

    # Verify TTL order
    early_ttl = in_memory_cache.ttl_dict["early_expire"]
    late_ttl = in_memory_cache.ttl_dict["late_expire"]
    assert early_ttl < late_ttl, "early_expire should have earlier expiration time"

    assert len(in_memory_cache.cache_dict) == 2

    # Add third item - should evict the one with earliest expiration time
    in_memory_cache.set_cache(key="new_item", value="value_3", ttl=300)

    assert len(in_memory_cache.cache_dict) == 2

    # Item with earliest expiration should be evicted
    assert "early_expire" not in in_memory_cache.cache_dict
    assert "early_expire" not in in_memory_cache.ttl_dict

    # Items with later expiration should remain
    assert "late_expire" in in_memory_cache.cache_dict
    assert "new_item" in in_memory_cache.cache_dict


def test_in_memory_cache_heap_size_staus_bounded():
    """
    Test that the expiration_heap does not grow unbounded when the same key is updated repeaatedly.
    """
    in_memory_cache = InMemoryCache(max_size_in_memory=10)

    for i in range(1_000):
        in_memory_cache.set_cache(key="hot_key", value=f"value_{i}", ttl=60)

    # Expiration heap should only have 1 entry
    assert len(in_memory_cache.expiration_heap) == 1


def test_in_memory_cache_heap_stays_bounded_when_a_live_key_is_rewritten_with_a_new_ttl() -> None:
    in_memory_cache = InMemoryCache(max_size_in_memory=10)
    in_memory_cache.set_cache(key="expires_first", value="pin", ttl=300)

    for i in range(1_000):
        in_memory_cache.delete_cache("hot_key")
        in_memory_cache.set_cache(key="hot_key", value=f"value_{i}", ttl=3600)

    assert in_memory_cache.get_cache("hot_key") == "value_999"
    assert len(in_memory_cache.expiration_heap) <= 2 * in_memory_cache.max_size_in_memory + 1


@pytest.fixture
def fine_grained_thread_switching() -> Iterator[None]:
    previous_interval: Final = sys.getswitchinterval()
    sys.setswitchinterval(1e-6)
    yield
    sys.setswitchinterval(previous_interval)


def _keys_until(deadline: float, rng: random.Random) -> Iterator[str]:
    while time.monotonic() < deadline:
        yield f"hot_{rng.randrange(5)}" if rng.random() < 0.7 else f"cold_{rng.randrange(40)}"


def _rewrite_keys_until(cache: InMemoryCache, deadline: float, seed: int) -> None:
    rng: Final = random.Random(seed)
    for key in _keys_until(deadline, rng):
        cache.delete_cache(key)
        cache.set_cache(key, "pin", ttl=rng.choice((300, 3600)))


@pytest.mark.usefixtures("fine_grained_thread_switching")
def test_in_memory_cache_survives_threads_rewriting_keys_in_a_full_cache() -> None:
    cache: Final = InMemoryCache(max_size_in_memory=5)
    deadline: Final = time.monotonic() + 0.5
    thread_count: Final = 8

    with ThreadPoolExecutor(max_workers=thread_count) as executor:
        futures: Final = tuple(
            executor.submit(_rewrite_keys_until, cache, deadline, seed) for seed in range(thread_count)
        )

    assert [future.exception() for future in futures] == [None] * thread_count
    keys: Final = (*(f"hot_{i}" for i in range(5)), *(f"cold_{i}" for i in range(40)))
    assert sum(cache.get_cache(key) is not None for key in keys) <= cache.max_size_in_memory


def test_in_memory_cache_stays_bounded_when_entries_have_no_usable_expiry() -> None:
    cache: Final = InMemoryCache(max_size_in_memory=10)
    for i in range(30):
        cache.set_cache(key=f"nan_{i}", value="response", ttl="NaN")
    cache.set_cache(key="fresh", value="response", ttl=60)

    keys: Final = (*(f"nan_{i}" for i in range(30)), "fresh")
    retained: Final = tuple(key for key in keys if cache.get_cache(key) is not None)
    assert len(retained) <= cache.max_size_in_memory
    assert cache.get_cache("fresh") == "response"


def test_in_memory_cache_prunes_expired_heap_entries_below_capacity():
    """
    Re-inserting expired keys below capacity should not grow expiration_heap
    without bound.
    """
    in_memory_cache = InMemoryCache(max_size_in_memory=200, default_ttl=1)

    for cycle in range(3):
        for i in range(5):
            in_memory_cache.set_cache(key=f"key_{i}", value=f"value_{cycle}_{i}", ttl=1)
        time.sleep(1.1)

    for i in range(5):
        in_memory_cache.set_cache(key=f"key_{i}", value=f"value_final_{i}", ttl=1)

    assert len(in_memory_cache.cache_dict) == 5
    assert len(in_memory_cache.ttl_dict) == 5
    assert len(in_memory_cache.expiration_heap) == 5


def test_in_memory_cache_injected_clock_controls_expiry_and_eviction() -> None:
    class Clock:
        now = 0.0

        def __call__(self) -> float:
            return self.now

    clock = Clock()
    cache = InMemoryCache(max_size_in_memory=2, default_ttl=60, clock=clock)
    cache.set_cache("first", "original", ttl=10)
    clock.now = 9.0
    cache.set_cache("second", "survivor")
    assert cache.get_cache("first") == "original"
    clock.now = 10.001
    assert cache.get_cache("first") is None
    cache.set_cache("third", "replacement")
    assert cache.get_cache("second") == "survivor"
    clock.now = 69.001
    cache.set_cache("fourth", "new")
    assert cache.get_cache("second") is None
    assert cache.get_cache("third") == "replacement"
    assert cache.get_cache("fourth") == "new"