litellm/tests/test_litellm_rust/cache/test_azure_blob.py

import asyncio
import json
import os
import time
import uuid
from collections.abc import Generator
from types import SimpleNamespace
from typing import Final, cast

import pytest
from azure.storage.blob import ContainerClient

from litellm.caching.azure_blob_cache import AzureBlobCache
from litellm.caching.caching import Cache
from litellm.rust_bridge import _native
from litellm.types.caching import LiteLLMCacheType
from tests.test_litellm_rust.support.cache import (
    CacheLookup,
    CacheTestHandle,
    CacheTestResolver,
    assert_native_runtime,
    completion_kwargs,
    request,
    require_rust,
)
from tests.test_litellm_rust.support.isolation import rebound

pytestmark: Final = pytest.mark.requires_rust_extension


@pytest.fixture
def azure_blob_facade() -> Generator[Cache]:
    account_url: Final = os.environ.get("AZURE_BLOB_CACHE_ACCOUNT_URL")
    if account_url is None:
        pytest.skip(
            "live Azure Blob parity needs AZURE_BLOB_CACHE_ACCOUNT_URL plus DefaultAzureCredential inputs in the environment"
        )
    facade: Final = Cache(
        type=LiteLLMCacheType.AZURE_BLOB,
        azure_account_url=account_url,
        azure_blob_container=f"litellm-parity-{uuid.uuid4().hex[:12]}",
    )
    backend: Final = facade.cache
    assert isinstance(backend, AzureBlobCache)
    try:
        yield facade
    finally:
        backend.container_client.delete_container()
        asyncio.run(backend.disconnect())


def azure_blob_handle(facade: Cache) -> _native._CacheTestHandle:
    backend: Final = facade.cache
    assert isinstance(backend, AzureBlobCache)
    return CacheTestHandle.azure_blob(
        backend.container_client.url.removesuffix(f"/{backend.container_client.container_name}"),
        backend.container_client.container_name,
    )


def test_azure_blob_facade_serves_natively_and_python_reads_the_same_blobs(azure_blob_facade: Cache) -> None:
    backend: Final = azure_blob_facade.cache
    assert isinstance(backend, AzureBlobCache)
    handle: Final = azure_blob_handle(azure_blob_facade)
    assert handle.backend == "azure-blob"
    account_url: Final = backend.container_client.url.removesuffix(f"/{backend.container_client.container_name}")
    with pytest.raises(TypeError, match="containers must match"):
        CacheTestHandle.azure_blob(account_url, f"{backend.container_client.container_name}-other")._bind_facade(
            azure_blob_facade
        )
    handle._bind_facade(azure_blob_facade)
    resolver: Final = CacheTestResolver(SimpleNamespace(cache=azure_blob_facade))
    native: Final = resolver.resolve()
    assert native.kind == "native"

    response: Final = {
        "choices": [{"text": "caf\u00e9 \u2603"}],
        "usage": {"total_tokens": 3},
        "flag": True,
        "empty": None,
    }
    native.store({**request("sync"), "ttl_seconds": 0.001}, response)
    native.store(request("sync"), {"choices": [{"text": "second"}]})
    time.sleep(0.01)
    stored: Final = json.loads(backend.container_client.download_blob("sync").readall())
    assert stored["response"] == response
    assert isinstance(stored["timestamp"], float)
    assert native.lookup(request("sync")) == response
    assert cast(CacheLookup, azure_blob_facade).get_cache(cache_key="sync") == response

    backend.set_cache("python", {"timestamp": time.time(), "response": response})
    backend.set_cache("legacy", "bare legacy value")
    backend.container_client.upload_blob("invalid", b"{not json", overwrite=True)
    assert native.lookup(request("python")) == response
    assert native.lookup(request("legacy")) == cast(CacheLookup, azure_blob_facade).get_cache(cache_key="legacy")
    assert native.lookup_batch([request("python"), request("missing"), request("invalid"), request("sync")]) == {
        "values": [response, None, None, response],
        "missing_indices": [1, 2],
    }

    with rebound(azure_blob_facade, "ttl", 12):
        assert resolver.resolve().kind == "python_callback"
    with rebound(backend, "container_client", ContainerClient.from_container_url(backend.container_client.url)):
        assert resolver.resolve().kind == "python_callback"

    def custom_get(*_args: object, **_kwargs: object) -> None:
        return None

    with rebound(backend, "get_cache", custom_get):
        assert resolver.resolve().kind == "python_callback"
    assert resolver.resolve().kind == "python_callback"
    assert cast(CacheLookup, azure_blob_facade).get_cache(cache_key="sync") == response

    class CustomBlobCache(AzureBlobCache):
        pass

    with rebound(azure_blob_facade, "cache", CustomBlobCache(account_url, backend.container_client.container_name)):
        assert resolver.resolve().kind == "python_callback"
        with pytest.raises(TypeError):
            azure_blob_handle(azure_blob_facade)._bind_facade(azure_blob_facade)


async def test_azure_blob_native_async_writes_overwrite_batch_and_flush_like_python(azure_blob_facade: Cache) -> None:
    backend: Final = azure_blob_facade.cache
    assert isinstance(backend, AzureBlobCache)
    azure_blob_handle(azure_blob_facade)._bind_facade(azure_blob_facade)
    binding: Final = CacheTestResolver(SimpleNamespace(cache=azure_blob_facade)).resolve()
    assert binding.kind == "native"
    ping: Final = cast(dict[str, object], await binding.ping())
    assert ping["status"] == "success", ping

    await binding.async_store(request("async"), {"value": 1})
    await binding.async_store({**request("async"), "ttl_seconds": 0.001}, {"value": 2})
    time.sleep(0.01)
    assert await binding.async_lookup(request("async")) == {"value": 2}
    assert await backend.async_get_cache("async") == json.loads(
        backend.container_client.download_blob("async").readall()
    )
    assert cast(CacheLookup, azure_blob_facade).get_cache(cache_key="async") == {"value": 2}

    await binding.async_store_batch([request("first"), request("second")], [{"value": 3}, {"value": 4}])
    assert await binding.async_lookup_batch([request("second"), request("missing"), request("first")]) == {
        "values": [{"value": 4}, None, {"value": 3}],
        "missing_indices": [1],
    }
    await binding.async_flush()
    assert [blob.name for blob in backend.container_client.list_blobs()] == []
    assert await binding.async_lookup(request("async")) is None


async def test_azure_blob_rust_required_rule_activates_natively(monkeypatch: pytest.MonkeyPatch) -> None:
    account_url: Final = os.environ.get("AZURE_BLOB_CACHE_ACCOUNT_URL")
    if account_url is None:
        pytest.skip(
            "live Azure Blob parity needs AZURE_BLOB_CACHE_ACCOUNT_URL plus DefaultAzureCredential inputs in the environment"
        )
    require_rust(monkeypatch, LiteLLMCacheType.AZURE_BLOB)
    facade: Final = Cache(
        type=LiteLLMCacheType.AZURE_BLOB,
        azure_account_url=account_url,
        azure_blob_container=f"litellm-parity-{uuid.uuid4().hex[:12]}",
    )
    backend: Final = facade.cache
    assert isinstance(backend, AzureBlobCache)
    try:
        assert_native_runtime(facade)
        kwargs: Final = completion_kwargs("azure")
        await facade.async_add_cache({"answer": "azure"}, **kwargs)
        assert await facade.async_get_cache(**kwargs) == {"answer": "azure"}
        assert backend.get_cache(facade.get_cache_key(**kwargs))["response"] == {"answer": "azure"}
    finally:
        backend.container_client.delete_container()
        await backend.disconnect()