litellm/tests/test_litellm_rust/cache/test_s3.py

import json
import time
from datetime import datetime
from types import SimpleNamespace
from typing import Final, cast
from unittest.mock import Mock

import boto3
import botocore.config
import pytest

from litellm.caching.caching import Cache
from litellm.caching.s3_cache import S3Cache
from litellm.rust_bridge import _native
from litellm.types.caching import LiteLLMCacheType
from tests.test_litellm_rust.support.cache import CacheTestResolver, activate_native, native_runtime, request
from tests.test_litellm_rust.support.isolation import rebound
from tests.test_litellm_rust.support.s3_stub import S3Stub

pytestmark: Final = pytest.mark.requires_rust_extension


def python_s3(url: str) -> S3Cache:
    return S3Cache(
        s3_bucket_name="cache-bucket",
        s3_region_name="us-east-1",
        s3_endpoint_url=url,
        s3_aws_access_key_id="key",
        s3_aws_secret_access_key="secret",
        s3_path="team",
    )


def s3_facade(url: str) -> Cache:
    return Cache(
        type=LiteLLMCacheType.S3,
        s3_bucket_name="cache-bucket",
        s3_region_name="us-east-1",
        s3_endpoint_url=url,
        s3_aws_access_key_id="key",
        s3_aws_secret_access_key="secret",
        s3_path="team",
    )


async def test_s3_reads_python_entries_and_writes_with_python_metadata(s3_stub: S3Stub) -> None:
    python_cache: Final = python_s3(s3_stub.url)
    response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}}
    python_cache.set_cache("sync:key", {"timestamp": time.time(), "response": response}, ttl=90)
    python_cache.set_cache("plain", {"timestamp": time.time(), "response": response})
    s3_stub.put_object("team/malformed", b"not a cache entry")
    s3_stub.put_object(
        "team/expired",
        json.dumps({"timestamp": time.time(), "response": response}).encode(),
        {"expires": "Thu, 01 Jan 1970 00:00:00 GMT"},
    )
    binding: Final = native_runtime(s3_facade(s3_stub.url))

    assert binding.lookup(request("sync:key")) == response
    assert await binding.async_lookup(request("plain")) == response
    assert binding.lookup(request("malformed")) is None
    assert binding.lookup(request("expired")) is None
    assert binding.lookup(request("absent")) is None

    binding.store({**request("native:key"), "ttl_seconds": 90.0}, response)
    await binding.async_store(request("no_ttl"), response)
    stored: Final = s3_stub.objects["team/native/key"]
    assert stored.headers["content-type"] == "application/json"
    assert stored.headers["content-language"] == "en"
    assert stored.headers["content-disposition"] == 'inline; filename="team/native/key.json"'
    assert stored.headers["cache-control"] == "immutable, max-age=90, s-maxage=90"
    expires: Final = cast(datetime, s3_stub.expires("team/native/key"))
    remaining: Final = (expires - datetime.now(expires.tzinfo)).total_seconds()
    assert 60 < remaining <= 91
    no_ttl: Final = s3_stub.objects["team/no_ttl"]
    assert no_ttl.headers["cache-control"] == "immutable, max-age=31536000, s-maxage=31536000"
    assert "expires" not in no_ttl.headers
    assert python_cache.get_cache("native:key")["response"] == response

    partial: Final = await binding.async_lookup_batch([request("native:key"), request("absent"), request("malformed")])
    assert partial == {"values": [response, None, None], "missing_indices": [1, 2]}


def test_selected_s3_runtime_declines_backend_mutation(s3_stub: S3Stub) -> None:
    facade: Final = Cache(
        type=LiteLLMCacheType.S3,
        s3_bucket_name="cache-bucket",
        s3_region_name="us-east-1",
        s3_endpoint_url=s3_stub.url,
        s3_aws_access_key_id="key",
        s3_aws_secret_access_key="secret",
        s3_path="team",
    )
    activate_native(facade)
    resolver: Final = CacheTestResolver(SimpleNamespace(cache=facade))
    binding: Final = resolver.resolve()
    assert binding.kind == "native"

    handler: Final = Mock()
    facade.cache.s3_client.meta.events.register("before-call.s3.*", handler)
    binding.store(request("native"), {"answer": 1})
    assert binding.lookup(request("native")) == {"answer": 1}
    assert handler.call_count == 0
    assert "team/native" in s3_stub.objects

    with rebound(facade.cache, "bucket_name", "other"):
        with pytest.raises(_native.RustBridgeDeclined):
            resolver.resolve()
    other_client: Final = boto3.client(
        "s3",
        region_name="us-east-1",
        endpoint_url=s3_stub.url,
        aws_access_key_id="key",
        aws_secret_access_key="secret",
    )
    with rebound(facade.cache, "s3_client", other_client):
        with pytest.raises(_native.RustBridgeDeclined):
            resolver.resolve()

    class CustomS3Cache(S3Cache):
        pass

    subclassed: Final = Cache(
        type=LiteLLMCacheType.S3,
        s3_bucket_name="cache-bucket",
        s3_region_name="us-east-1",
        s3_endpoint_url=s3_stub.url,
        s3_aws_access_key_id="key",
        s3_aws_secret_access_key="secret",
        s3_path="team",
    )
    subclassed.cache = CustomS3Cache(
        s3_bucket_name="cache-bucket",
        s3_region_name="us-east-1",
        s3_endpoint_url=s3_stub.url,
        s3_aws_access_key_id="key",
        s3_aws_secret_access_key="secret",
        s3_path="team",
    )
    assert CacheTestResolver(SimpleNamespace(cache=subclassed)).resolve().kind == "python_callback"


def test_s3_facade_rejects_configurations_that_require_python(s3_stub: S3Stub) -> None:
    unverified: Final = Cache(
        type=LiteLLMCacheType.S3,
        s3_bucket_name="cache-bucket",
        s3_region_name="us-east-1",
        s3_endpoint_url="https://s3.example.test",
        s3_aws_access_key_id="key",
        s3_aws_secret_access_key="secret",
        s3_path="team",
        s3_verify=False,
    )
    with pytest.raises(_native.RustBridgeDeclined, match="requires Python"):
        native_runtime(unverified)
    proxied: Final = Cache(
        type=LiteLLMCacheType.S3,
        s3_bucket_name="cache-bucket",
        s3_region_name="us-east-1",
        s3_endpoint_url=s3_stub.url,
        s3_aws_access_key_id="key",
        s3_aws_secret_access_key="secret",
        s3_path="team",
        s3_config=botocore.config.Config(proxies={"https": "http://proxy.test"}),
    )
    with pytest.raises(_native.RustBridgeDeclined, match="requires Python"):
        native_runtime(proxied)