From a6d1497932e257a809fdd4647d85a70039376816 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:38:20 +0000 Subject: [PATCH] test(rust): add GCS native cache parity fixtures Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 1 - litellm-rust/crates/cache-gcs/Cargo.toml | 1 - litellm-rust/crates/cache-gcs/src/cache.rs | 37 +-- litellm-rust/crates/cache-gcs/tests/cache.rs | 34 +++ tests/test_litellm_rust/support/fake_gcs.py | 152 +++++++++++++ tests/test_litellm_rust/test_cache.py | 228 +++++++++++++++++++ 6 files changed, 420 insertions(+), 33 deletions(-) create mode 100644 tests/test_litellm_rust/support/fake_gcs.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 99a6133b4f1..ae28553801d 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2476,7 +2476,6 @@ dependencies = [ "reqwest 0.12.28", "serde_json", "tokio", - "url", "wiremock", ] diff --git a/litellm-rust/crates/cache-gcs/Cargo.toml b/litellm-rust/crates/cache-gcs/Cargo.toml index 3890e60843c..4ec60bcfa3b 100644 --- a/litellm-rust/crates/cache-gcs/Cargo.toml +++ b/litellm-rust/crates/cache-gcs/Cargo.toml @@ -13,7 +13,6 @@ litellm-cache.workspace = true percent-encoding.workspace = true reqwest.workspace = true tokio.workspace = true -url.workspace = true [dev-dependencies] serde_json.workspace = true diff --git a/litellm-rust/crates/cache-gcs/src/cache.rs b/litellm-rust/crates/cache-gcs/src/cache.rs index 1097f3252b1..65282ac99d5 100644 --- a/litellm-rust/crates/cache-gcs/src/cache.rs +++ b/litellm-rust/crates/cache-gcs/src/cache.rs @@ -5,43 +5,18 @@ use litellm_cache::{ BaseCache, BatchCache, BatchEntry, CacheCodec, CacheConnectionResult, Error, ExactCacheContext, FlushCache, }; -use percent_encoding::{AsciiSet, CONTROLS, percent_encode}; +use percent_encoding::{AsciiSet, NON_ALPHANUMERIC, percent_encode}; use reqwest::Client; use crate::{GcpTokenSource, TokenSource}; pub const DEFAULT_ENDPOINT: &str = "https://storage.googleapis.com"; -const OBJECT_NAME_ENCODE_SET: &AsciiSet = &CONTROLS - .add(b' ') - .add(b'!') - .add(b'"') - .add(b'#') - .add(b'$') - .add(b'%') - .add(b'&') - .add(b'\'') - .add(b'(') - .add(b')') - .add(b'*') - .add(b'+') - .add(b',') - .add(b'/') - .add(b':') - .add(b';') - .add(b'<') - .add(b'=') - .add(b'>') - .add(b'?') - .add(b'@') - .add(b'[') - .add(b'\\') - .add(b']') - .add(b'^') - .add(b'`') - .add(b'{') - .add(b'|') - .add(b'}'); +const OBJECT_NAME_ENCODE_SET: &AsciiSet = &NON_ALPHANUMERIC + .remove(b'-') + .remove(b'_') + .remove(b'.') + .remove(b'~'); pub fn key_prefix(gcs_path: Option<&str>) -> String { match gcs_path { diff --git a/litellm-rust/crates/cache-gcs/tests/cache.rs b/litellm-rust/crates/cache-gcs/tests/cache.rs index 79e4ed100f8..45eecf01cec 100644 --- a/litellm-rust/crates/cache-gcs/tests/cache.rs +++ b/litellm-rust/crates/cache-gcs/tests/cache.rs @@ -117,6 +117,40 @@ fn key_prefix_normalizes_paths() { assert_eq!(key_prefix(Some("")), ""); } +#[tokio::test] +async fn object_names_use_python_quote_encoding() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/upload/storage/v1/b/bucket/o")) + .and(query_param("uploadType", "media")) + .respond_with(ResponseTemplate::new(200)) + .expect(2) + .mount(&server) + .await; + let cache = cache(&server, Some("p/")); + cache + .set_cache( + "a~b-c_d.e/f g%h", + json!({"value": "punctuation"}), + &ExactCacheContext::default(), + ) + .unwrap(); + cache + .set_cache( + "ключ", + json!({"value": "utf8"}), + &ExactCacheContext::default(), + ) + .unwrap(); + let requests = server.received_requests().await.unwrap(); + let queries: Vec<_> = requests + .iter() + .filter_map(|request| request.url.query()) + .collect(); + assert!(queries.contains(&"uploadType=media&name=p%2Fa~b-c_d.e%2Ff%20g%25h")); + assert!(queries.contains(&"uploadType=media&name=p%2F%D0%BA%D0%BB%D1%8E%D1%87")); +} + #[tokio::test] async fn ignores_ttl_and_writes_pipeline_concurrently() { let server = MockServer::start().await; diff --git a/tests/test_litellm_rust/support/fake_gcs.py b/tests/test_litellm_rust/support/fake_gcs.py new file mode 100644 index 00000000000..67eb61798b9 --- /dev/null +++ b/tests/test_litellm_rust/support/fake_gcs.py @@ -0,0 +1,152 @@ +from __future__ import annotations + +import json +import threading +from collections.abc import Mapping +from dataclasses import dataclass +from functools import partial +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from socket import socket +from types import MappingProxyType +from typing import Final, cast +from urllib.parse import unquote, urlsplit + + +@dataclass(frozen=True, slots=True) +class RecordedRequest: + method: str + path: str + query: str + headers: Mapping[str, str] + body: bytes + + +class _FakeGcsHandler(BaseHTTPRequestHandler): + def __init__( + self, + request: socket | tuple[bytes, socket], + client_address: tuple[str, int], + server: ThreadingHTTPServer, + *, + fake: FakeGcs, + ) -> None: + self._fake: Final = fake + super().__init__(request, client_address, server) + + def _handle(self) -> None: + parsed: Final = urlsplit(self.path) + content_length: Final = int(self.headers.get("Content-Length", "0")) + body: Final = self.rfile.read(content_length) if content_length else b"" + headers: Final = MappingProxyType( + {name.title(): value for name, value in self.headers.items()} + ) + self._fake.record( + RecordedRequest( + method=self.command, + path=parsed.path, + query=parsed.query, + headers=headers, + body=body, + ) + ) + if self.headers.get("Authorization") != f"Bearer {self._fake.token}": + self._send_json(401, {"error": "unauthorized"}) + return + + upload_prefix: Final = "/upload/storage/v1/b/" + download_prefix: Final = "/storage/v1/b/" + if parsed.path.startswith(upload_prefix) and parsed.path.endswith("/o"): + self._upload(parsed.path[len(upload_prefix) : -2], parsed.query, body) + return + if parsed.path.startswith(download_prefix): + self._download(parsed.path[len(download_prefix) :], parsed.query) + return + self._send_json(404, {"error": "not found"}) + + def _upload(self, path: str, query: str, body: bytes) -> None: + values: Final = { + unquote(pair.partition("=")[0]): unquote(pair.partition("=")[2]) + for pair in query.split("&") + if pair + } + if not path or values.get("uploadType") != "media" or "name" not in values: + self._send_json(404, {"error": "not found"}) + return + self._fake.put_object(path, values["name"], body) + self._send_json(200, {"name": values["name"], "bucket": path}) + + def _download(self, path: str, query: str) -> None: + bucket, separator, encoded_name = path.partition("/o/") + if not separator or query != "alt=media": + self._send_json(404, {"error": "not found"}) + return + name: Final = unquote(encoded_name) + if name.endswith("/server-error") or name == "server-error": + self._send_json(500, {"error": "server error"}) + return + body: Final = self._fake.get_object(bucket, name) + if body is None: + self._send_json(404, {"error": "not found"}) + return + self._send(200, body, "application/octet-stream") + + def _send_json(self, status: int, value: object) -> None: + payload: Final = json.dumps(value).encode() + self._send(status, payload, "application/json") + + def _send(self, status: int, body: bytes, content_type: str) -> None: + self.send_response(status) + self.send_header("Content-Type", content_type) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, format: str, *args: object) -> None: + pass + + do_GET = _handle + do_POST = _handle + + +class FakeGcs: + def __init__(self) -> None: + self._objects: dict[tuple[str, str], bytes] = {} # mutable-ok: fake object store + self._requests: list[RecordedRequest] = [] # mutable-ok: recorded request history + self._server = ThreadingHTTPServer( + ("127.0.0.1", 0), + partial(_FakeGcsHandler, fake=self), + ) + self._worker = threading.Thread(target=self._server.serve_forever, daemon=True) + self._worker.start() + self.token: Final = "test-token" + + @property + def url(self) -> str: + address: Final = cast(tuple[str, int], self._server.server_address) + host, port = address + return f"http://{host}:{port}" + + @property + def objects(self) -> Mapping[tuple[str, str], bytes]: + return MappingProxyType(self._objects) + + @property + def requests(self) -> tuple[RecordedRequest, ...]: + return tuple(self._requests) + + def put(self, bucket: str, name: str, body: bytes) -> None: + self.put_object(bucket, name, body) + + def close(self) -> None: + self._server.shutdown() + self._server.server_close() + self._worker.join(timeout=5) + + def record(self, request: RecordedRequest) -> None: + self._requests.append(request) + + def put_object(self, bucket: str, name: str, body: bytes) -> None: + self._objects[(bucket, name)] = body + + def get_object(self, bucket: str, name: str) -> bytes | None: + return self._objects.get((bucket, name)) diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index c35cb1a20fb..cc409ebe49c 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -16,9 +16,11 @@ import redis import litellm from litellm.caching.caching import Cache, disable_cache, enable_cache, update_cache +from litellm.caching.gcs_cache import GCSCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.rust_bridge import _native from litellm.types.caching import LiteLLMCacheType +from tests.test_litellm_rust.support.fake_gcs import FakeGcs from tests.test_litellm_rust.support.isolation import rebound pytestmark: Final = pytest.mark.requires_rust_extension @@ -26,6 +28,7 @@ pytestmark: Final = pytest.mark.requires_rust_extension class CacheLookup(Protocol): def get_cache(self, **kwargs: object) -> object: ... + def flush_cache(self) -> object: ... def request(key: str = "key") -> dict[str, object]: @@ -45,6 +48,15 @@ def redis_url() -> Generator[str]: worker.join(timeout=5) +@pytest.fixture +def fake_gcs() -> Generator[FakeGcs]: + server: Final = FakeGcs() + try: + yield server + finally: + server.close() + + def test_existing_constructor_and_global_are_unchanged() -> None: facade: Final = Cache(type=LiteLLMCacheType.LOCAL) assert type(facade.cache) is InMemoryCache @@ -393,3 +405,219 @@ async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None: assert client.get("second") is not None await facade.cache.disconnect() client.close() + + +async def test_gcs_reads_python_entries_and_writes_python_compatible_objects( + fake_gcs: FakeGcs, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False) + monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) + response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}, "flag": True, "empty": None} + fake_gcs.put( + "bucket", + "cache/sync", + json.dumps({"timestamp": time.time(), "response": json.dumps(response)}).encode(), + ) + fake_gcs.put("bucket", "cache/async", json.dumps({"timestamp": time.time(), "response": response}).encode()) + fake_gcs.put("bucket", "cache/raw", json.dumps(response).encode()) + fake_gcs.put("bucket", "cache/invalid", b"not a cache entry") + binding: Final = _native._CacheTestResolver( + SimpleNamespace( + cache=_native._CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + ) + ).resolve() + + assert binding.lookup(request("sync")) == response + assert await binding.async_lookup(request("async")) == response + assert binding.lookup(request("raw")) == response + assert await binding.async_lookup(request("invalid")) is None + assert binding.lookup(request("missing")) is None + + await binding.async_store({**request("native"), "ttl_seconds": 12.0}, response) + stored: Final = fake_gcs.objects[("bucket", "cache/native")] + stored_value: Final = cast(dict[str, object], json.loads(stored)) + assert stored_value["response"] == response + assert isinstance(stored_value["timestamp"], float) + upload: Final = next(item for item in fake_gcs.requests if item.method == "POST") + assert upload.path == "/upload/storage/v1/b/bucket/o" + assert upload.query == "uploadType=media&name=cache%2Fnative" + assert upload.headers["Authorization"] == f"Bearer {fake_gcs.token}" + assert upload.headers["Content-Type"] == "application/json" + upload_text: Final = f"{upload.path}?{upload.query}{upload.headers}" + assert "ttl" not in upload_text.lower() + assert "expiry" not in upload_text.lower() + download: Final = next(item for item in fake_gcs.requests if item.path.endswith("/cache%2Fsync")) + assert download.path == "/storage/v1/b/bucket/o/cache%2Fsync" + assert download.query == "alt=media" + + binding.store(request("sync2"), response) + assert binding.lookup(request("sync2")) == response + assert GCSCache(bucket_name="bucket", gcs_path="cache").key_prefix == "cache/" + assert GCSCache(bucket_name="bucket", gcs_path="cache/").key_prefix == "cache/" + assert GCSCache(bucket_name="bucket").key_prefix == "" + + +async def test_gcs_batch_lookup_preserves_order_and_treats_malformed_entries_as_misses(fake_gcs: FakeGcs) -> None: + fake_gcs.put("bucket", "cache/hit", json.dumps({"timestamp": time.time(), "response": {"value": 1}}).encode()) + fake_gcs.put("bucket", "cache/invalid", b"not a cache entry") + binding: Final = _native._CacheTestResolver( + SimpleNamespace( + cache=_native._CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + ) + ).resolve() + requests: Final = [request("hit"), request("missing"), request("invalid")] + expected: Final = {"values": [{"value": 1}, None, None], "missing_indices": [1, 2]} + + assert await binding.async_lookup_batch(requests) == expected + assert binding.lookup_batch(requests) == expected + await binding.async_store_batch([request("first"), request("second")], [{"value": 1}, {"value": 2}]) + assert ("bucket", "cache/first") in fake_gcs.objects + assert ("bucket", "cache/second") in fake_gcs.objects + + +async def test_gcs_facade_binds_only_exact_matching_configuration( + fake_gcs: FakeGcs, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False) + monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) + monkeypatch.setenv("GOOGLE_APPLICATION_CREDENTIALS", "/nonexistent") + facade: Final = Cache(type=LiteLLMCacheType.GCS, gcs_bucket_name="bucket", gcs_path="cache/") + assert type(facade.cache) is GCSCache + + mismatched_bucket: Final = _native._CacheTestHandle.gcs( + "other", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + with pytest.raises(TypeError, match="buckets must match"): + mismatched_bucket._bind_facade(facade) + mismatched_prefix: Final = _native._CacheTestHandle.gcs( + "bucket", + gcs_path="x", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + with pytest.raises(TypeError, match="key prefixes must match"): + mismatched_prefix._bind_facade(facade) + mismatched_credentials: Final = _native._CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + path_service_account="sa.json", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + with pytest.raises(TypeError, match="credentials must match"): + mismatched_credentials._bind_facade(facade) + with pytest.raises(TypeError, match="types must match"): + _native._CacheTestHandle.memory()._bind_facade(facade) + + matching: Final = _native._CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + matching._bind_facade(facade) + resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) + binding: Final = resolver.resolve() + assert binding.kind == "native" + await binding.async_store(request("native"), {"value": "native"}) + assert await binding.async_lookup(request("native")) == {"value": "native"} + assert cast(CacheLookup, facade).get_cache(cache_key="native") is None + + with rebound(facade.cache, "bucket_name", "other"): + assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "key_prefix", "x/"): + assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "path_service_account", "sa.json"): + assert resolver.resolve().kind == "python_callback" + def no_get_cache(*args: object, **kwargs: object) -> None: + return None + + with rebound(facade.cache, "get_cache", no_get_cache): + assert resolver.resolve().kind == "python_callback" + with rebound(facade, "ttl", 12): + assert resolver.resolve().kind == "python_callback" + + class CustomGcs(GCSCache): + pass + + with rebound(facade, "cache", CustomGcs(bucket_name="bucket", gcs_path="cache/")): + assert resolver.resolve().kind == "python_callback" + custom_facade: Final = Cache(type=LiteLLMCacheType.GCS, gcs_bucket_name="bucket", gcs_path="cache/") + with rebound(custom_facade, "cache", CustomGcs(bucket_name="bucket", gcs_path="cache/")): + with pytest.raises(TypeError, match="types must match"): + matching._bind_facade(custom_facade) + + missing_bucket: Final = Cache(type=LiteLLMCacheType.GCS) + with pytest.raises(TypeError, match="requires a configured bucket name"): + matching._bind_facade(missing_bucket) + + +async def test_gcs_flush_is_a_no_op_and_ping_is_not_implemented( + fake_gcs: FakeGcs, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False) + monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) + binding: Final = _native._CacheTestResolver( + SimpleNamespace( + cache=_native._CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + ) + ).resolve() + await binding.async_store(request("key"), {"value": "stored"}) + await binding.async_flush() + assert ("bucket", "cache/key") in fake_gcs.objects + assert await binding.async_lookup(request("key")) == {"value": "stored"} + with pytest.raises(NotImplementedError): + await binding.ping() + + facade: Final = Cache(type=LiteLLMCacheType.GCS, gcs_bucket_name="bucket", gcs_path="cache/") + with pytest.raises(AttributeError): + await facade.ping() + assert cast(CacheLookup, facade.cache).flush_cache() is None + + +async def test_gcs_unauthorized_and_server_errors_surface_as_runtime_errors(fake_gcs: FakeGcs) -> None: + wrong_token: Final = _native._CacheTestResolver( + SimpleNamespace( + cache=_native._CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token="wrong-token", + ) + ) + ).resolve() + with pytest.raises(RuntimeError): + wrong_token.lookup(request("missing")) + assert not fake_gcs.objects + + binding: Final = _native._CacheTestResolver( + SimpleNamespace( + cache=_native._CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + ) + ).resolve() + with pytest.raises(RuntimeError): + binding.lookup(request("server-error")) + assert binding.lookup(request("missing")) is None