Spaces:
Configuration error
Configuration error
File size: 6,836 Bytes
447ebeb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 |
import json
import os
import sys
import pytest
from fastapi.testclient import TestClient
sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
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
assert "traceback" in error_details
# Verify specific error message
assert "invalid username-password pair" in error_details["message"]
def test_cache_ping_no_cache_initialized():
"""Test cache ping when no cache is initialized"""
# Set cache to None
original_cache = litellm.cache
litellm.cache = None
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))
assert "error" in data
error = data["error"]
# Verify error contains all expected fields
assert "message" in error
error_details = json.loads(error["message"])
assert "Cache not initialized. litellm.cache is None" in error_details["message"]
# Restore original cache
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)
|