mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
6a0d03914c
* test: drop the cwd-relative sys.path.insert calls from the test suite
TQ003 stands at 1,077 across 1,058 files, and 1,015 of them are the same shape:
sys.path.insert(0, os.path.abspath("../..")) and its deeper siblings. The
argument resolves against the working directory rather than the file, so from
the repo root, where every job runs pytest, it inserts the directory two levels
above the checkout. It has never pointed at litellm. The package is installed
into the environment anyway, which is what actually makes the import work, and
what the rule's message has said all along.
Removing them leaves 1,634 imports of sys and os with no remaining reference,
and those go too, except where another test module imports the name back out of
the file. The rest of TQ003 is 62 call sites that resolve against __file__ or a
variable, which are a different question and are left alone.
Collection is identical either way: 45,871 tests and the same 51 pre-existing
collection errors before and after, and ruff reports no new undefined name.
* test: drop the duplicate imports the sys.path sweep exposed to F811
* test(pre-call-utils): restore the os import the new bedrock tests need
336 lines
11 KiB
Python
336 lines
11 KiB
Python
import json
|
|
|
|
import pytest
|
|
from fastapi.testclient import TestClient
|
|
|
|
|
|
|
|
import litellm
|
|
from litellm.caching import RedisCache
|
|
from litellm.proxy.proxy_server import app
|
|
|
|
client = TestClient(app)
|
|
|
|
|
|
# Mock successful Redis connection
|
|
@pytest.fixture
|
|
def mock_redis_success(mocker):
|
|
async def mock_ping():
|
|
return True
|
|
|
|
async def mock_add_cache(*args, **kwargs):
|
|
return None
|
|
|
|
mock_cache = mocker.MagicMock()
|
|
mock_cache.type = "redis"
|
|
mock_cache.ping = mock_ping
|
|
mock_cache.async_add_cache = mock_add_cache
|
|
mock_cache.cache = RedisCache(
|
|
host="localhost",
|
|
port=6379,
|
|
password="hello",
|
|
)
|
|
|
|
mocker.patch.object(litellm, "cache", mock_cache)
|
|
return mock_cache
|
|
|
|
|
|
# Mock failed Redis connection
|
|
@pytest.fixture
|
|
def mock_redis_failure(mocker):
|
|
async def mock_ping():
|
|
raise Exception("invalid username-password pair")
|
|
|
|
mock_cache = mocker.MagicMock()
|
|
mock_cache.type = "redis"
|
|
mock_cache.ping = mock_ping
|
|
|
|
mocker.patch.object(litellm, "cache", mock_cache)
|
|
return mock_cache
|
|
|
|
|
|
def test_cache_ping_success(mock_redis_success):
|
|
"""Test successful cache ping with regular response"""
|
|
response = client.get("/cache/ping", headers={"Authorization": "Bearer sk-1234"})
|
|
assert response.status_code == 200
|
|
|
|
data = response.json()
|
|
assert data["status"] == "healthy"
|
|
assert data["cache_type"] == "redis"
|
|
assert data["ping_response"] is True
|
|
assert data["set_cache_response"] == "success"
|
|
|
|
|
|
def test_cache_ping_with_complex_objects(mock_redis_success, mocker):
|
|
"""Test cache ping with non-standard serializable objects"""
|
|
|
|
# Mock complex objects in the cache parameters
|
|
class ComplexObject:
|
|
def __str__(self):
|
|
return "complex_object"
|
|
|
|
mock_redis_success.cache.complex_attr = ComplexObject()
|
|
mock_redis_success.cache.datetime_attr = mocker.MagicMock()
|
|
|
|
response = client.get("/cache/ping", headers={"Authorization": "Bearer sk-1234"})
|
|
assert response.status_code == 200
|
|
|
|
# Verify response is JSON serializable
|
|
data = response.json()
|
|
print("data=", json.dumps(data, indent=4))
|
|
assert data["status"] == "healthy"
|
|
assert "litellm_cache_params" in data
|
|
|
|
# Verify complex objects were converted to strings
|
|
cache_params = json.loads(data["litellm_cache_params"])
|
|
assert isinstance(cache_params, dict)
|
|
|
|
|
|
def test_cache_ping_with_circular_reference(mock_redis_success):
|
|
"""Test cache ping with circular reference in cache parameters"""
|
|
# Create circular reference
|
|
circular_dict = {}
|
|
circular_dict["self"] = circular_dict
|
|
mock_redis_success.cache.circular_ref = circular_dict
|
|
|
|
response = client.get("/cache/ping", headers={"Authorization": "Bearer sk-1234"})
|
|
assert response.status_code == 200
|
|
|
|
# Verify response is still JSON serializable
|
|
data = response.json()
|
|
assert data["status"] == "healthy"
|
|
|
|
|
|
def test_cache_ping_failure(mock_redis_failure):
|
|
"""Test cache ping failure with expected error fields"""
|
|
response = client.get("/cache/ping", headers={"Authorization": "Bearer sk-1234"})
|
|
assert response.status_code == 503
|
|
|
|
data = response.json()
|
|
print("data=", json.dumps(data, indent=4, default=str))
|
|
|
|
assert "error" in data
|
|
error = data["error"]
|
|
|
|
# Verify error contains all expected fields
|
|
assert "message" in error
|
|
error_details = json.loads(error["message"])
|
|
assert "message" in error_details
|
|
assert "litellm_cache_params" in error_details
|
|
assert "health_check_cache_params" in error_details
|
|
|
|
# Verify generic static message (exception text must not leak to clients)
|
|
assert error_details["message"] == "Service Unhealthy"
|
|
|
|
|
|
def test_cache_ping_failure_does_not_expose_traceback(mock_redis_failure):
|
|
"""CWE-209: Stack trace and exception text must not appear in the HTTP 503 response body."""
|
|
response = client.get("/cache/ping", headers={"Authorization": "Bearer sk-1234"})
|
|
assert response.status_code == 503
|
|
|
|
data = response.json()
|
|
error = data.get("error", {})
|
|
raw_body = json.dumps(data)
|
|
|
|
# The word "traceback" (case-insensitive) must not appear anywhere in the response
|
|
assert (
|
|
"traceback" not in raw_body.lower()
|
|
), "CWE-209: Python traceback exposed in HTTP 503 response body"
|
|
# Internal frame paths should not leak either
|
|
assert (
|
|
'File "' not in raw_body
|
|
), "CWE-209: Python stack frame paths exposed in HTTP 503 response body"
|
|
# Exception text (e.g. Redis hostnames/IPs) must not leak either
|
|
assert (
|
|
"invalid username-password pair" not in raw_body
|
|
), "CWE-209: Exception message text exposed in HTTP 503 response body"
|
|
|
|
# The error message should be the safe static string
|
|
error_details = json.loads(error["message"])
|
|
assert error_details["message"] == "Service Unhealthy"
|
|
|
|
|
|
def test_cache_ping_no_cache_initialized():
|
|
"""Test cache ping when no cache is initialized returns 503 with ProxyException envelope.
|
|
|
|
Verifies the exact response structure so that regressions in the error format
|
|
(e.g. message moving to a different field, or extra internal details leaking)
|
|
are caught immediately.
|
|
"""
|
|
original_cache = litellm.cache
|
|
litellm.cache = None
|
|
|
|
try:
|
|
response = client.get(
|
|
"/cache/ping", headers={"Authorization": "Bearer sk-1234"}
|
|
)
|
|
assert response.status_code == 503
|
|
|
|
data = response.json()
|
|
print("response data=", json.dumps(data, indent=4))
|
|
# ProxyException is serialised as {"error": {"message": "...", "type": ..., ...}}
|
|
assert "error" in data
|
|
error_details = json.loads(data["error"]["message"])
|
|
assert (
|
|
error_details["message"] == "Cache not initialized. litellm.cache is None"
|
|
)
|
|
finally:
|
|
litellm.cache = original_cache
|
|
|
|
|
|
def test_cache_ping_no_cache_does_not_expose_internals():
|
|
"""CWE-209: No-cache 503 must use the ProxyException envelope with no internal details.
|
|
|
|
The null-cache path raises ProxyException directly (not HTTPException), so the
|
|
response is {"error": {"message": "...", ...}} — same envelope as other 503s from
|
|
this endpoint — with no tracebacks, source paths, or extra fields leaking.
|
|
"""
|
|
original_cache = litellm.cache
|
|
litellm.cache = None
|
|
|
|
try:
|
|
response = client.get(
|
|
"/cache/ping", headers={"Authorization": "Bearer sk-1234"}
|
|
)
|
|
assert response.status_code == 503
|
|
|
|
raw_body = response.text
|
|
# No Python traceback or source-file paths must appear in the response
|
|
assert "traceback" not in raw_body.lower(), (
|
|
"CWE-209: Python traceback exposed in /cache/ping no-cache response"
|
|
)
|
|
assert 'File "' not in raw_body, (
|
|
"CWE-209: Python stack frame paths exposed in /cache/ping no-cache response"
|
|
)
|
|
|
|
data = response.json()
|
|
# Response must use the ProxyException envelope
|
|
assert "error" in data, f"Expected ProxyException envelope, got: {data}"
|
|
error_details = json.loads(data["error"]["message"])
|
|
assert (
|
|
error_details["message"] == "Cache not initialized. litellm.cache is None"
|
|
)
|
|
finally:
|
|
litellm.cache = original_cache
|
|
|
|
|
|
def test_cache_ping_health_check_includes_only_cache_attributes(mock_redis_success):
|
|
"""
|
|
Ensure that the /cache/ping endpoint only pulls HealthCheckCacheParams from litellm.cache.cache,
|
|
and not from other attributes on litellm.cache.
|
|
"""
|
|
# Add an unrelated field directly to the cache mock; it should NOT appear in health_check_cache_params
|
|
mock_redis_success.some_unrelated_field = "should-not-appear-in-health-check"
|
|
|
|
# Add a field on the underlying `cache` object that SHOULD appear
|
|
mock_redis_success.cache.redis_kwargs = {"host": "localhost", "port": 6379}
|
|
|
|
response = client.get("/cache/ping", headers={"Authorization": "Bearer sk-1234"})
|
|
assert (
|
|
response.status_code == 200
|
|
), f"Unexpected status code: {response.status_code}"
|
|
|
|
data = response.json()
|
|
print("/cache/ping response data=", json.dumps(data, indent=4))
|
|
health_check_cache_params = data.get("health_check_cache_params", {})
|
|
# The unrelated field we attached at the top-level of litellm.cache should *not* be present
|
|
assert (
|
|
"some_unrelated_field" not in health_check_cache_params
|
|
), "Found an unexpected field from the mock_redis_success object in health_check_cache_params"
|
|
|
|
# The field we attached to 'mock_redis_success.cache' should be present and correctly reported
|
|
assert (
|
|
"redis_kwargs" in health_check_cache_params
|
|
), "Expected field on `litellm.cache.cache` was not found in health_check_cache_params"
|
|
assert health_check_cache_params["redis_kwargs"] == {
|
|
"host": "localhost",
|
|
"port": 6379,
|
|
}
|
|
|
|
|
|
def test_cache_ping_with_redis_version_float(mock_redis_success):
|
|
"""Test cache ping works when redis_version is a float"""
|
|
# Set redis_version as a float
|
|
mock_redis_success.cache.redis_version = 7.2
|
|
|
|
response = client.get("/cache/ping", headers={"Authorization": "Bearer sk-1234"})
|
|
assert response.status_code == 200
|
|
|
|
data = response.json()
|
|
print("data=", json.dumps(data, indent=4))
|
|
assert data["status"] == "healthy"
|
|
assert data["cache_type"] == "redis"
|
|
|
|
cache_params = data["health_check_cache_params"]
|
|
assert isinstance(cache_params, dict)
|
|
assert isinstance(cache_params.get("redis_version"), float)
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_redis_client_list_restricted(mocker):
|
|
"""Mock Redis cache where CLIENT LIST is restricted (like GCP Redis)"""
|
|
|
|
def mock_client_list():
|
|
raise Exception("ERR unknown command 'CLIENT'")
|
|
|
|
def mock_info():
|
|
return {
|
|
"redis_version": "6.2.7",
|
|
"used_memory": "1000000",
|
|
"connected_clients": "5",
|
|
"keyspace_hits": "1000",
|
|
"keyspace_misses": "100",
|
|
}
|
|
|
|
mock_cache = mocker.MagicMock()
|
|
mock_cache.type = "redis"
|
|
mock_cache.cache = RedisCache(host="localhost", port=6379, password="hello")
|
|
mock_cache.cache.client_list = mock_client_list
|
|
mock_cache.cache.info = mock_info
|
|
|
|
mocker.patch.object(litellm, "cache", mock_cache)
|
|
return mock_cache
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_redis_client_list_success(mocker):
|
|
"""Mock Redis cache where CLIENT LIST works normally"""
|
|
|
|
def mock_client_list():
|
|
return [
|
|
{"id": "1", "addr": "127.0.0.1:54321", "name": "client1"},
|
|
{"id": "2", "addr": "127.0.0.1:54322", "name": "client2"},
|
|
]
|
|
|
|
def mock_info():
|
|
return {
|
|
"redis_version": "6.2.7",
|
|
"used_memory": "1000000",
|
|
"connected_clients": "2",
|
|
}
|
|
|
|
mock_cache = mocker.MagicMock()
|
|
mock_cache.type = "redis"
|
|
mock_cache.cache = RedisCache(host="localhost", port=6379, password="hello")
|
|
mock_cache.cache.client_list = mock_client_list
|
|
mock_cache.cache.info = mock_info
|
|
|
|
mocker.patch.object(litellm, "cache", mock_cache)
|
|
return mock_cache
|
|
|
|
|
|
def test_cache_redis_info_no_cache():
|
|
"""Test /cache/redis/info when no cache is initialized"""
|
|
original_cache = litellm.cache
|
|
litellm.cache = None
|
|
|
|
response = client.get(
|
|
"/cache/redis/info", headers={"Authorization": "Bearer sk-1234"}
|
|
)
|
|
assert response.status_code == 503
|
|
|
|
data = response.json()
|
|
assert "Cache not initialized" in data["detail"]
|
|
|
|
# Restore original cache
|
|
litellm.cache = original_cache
|