scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

fix(tests): restore litellm.model_cost after reload endpoint test

Browse source 
test_reload_model_cost_map_admin_access calls the /reload/model_cost_map
HTTP endpoint with get_model_cost_map mocked to return a single-entry
dict. The endpoint handler does a direct module-level assignment
(litellm.model_cost = new_model_cost_map) which persists after the
patch context manager exits, stripping all models except gpt-3.5-turbo
from the in-memory cost map and causing subsequent tests that rely on
models like gemini-1.5-flash, multimodalembedding@001, and gpt-4o to
fail with "model not mapped" errors or zero-cost spend payloads.

Fix: save litellm.model_cost before the test and restore it (along with
invalidating the case-insensitive lookup cache) in a finally block.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Julio Quinteros Pro 2026-02-18 18:20:48 -03:00
parent e8ab773ac4
commit 1c0f4302f8
1 changed files with 32 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

54
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -25,6 +25,7 @@ sys.path.insert(
import litellm
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.proxy_server import app, initialize
from litellm.utils import _invalidate_model_cost_lowercase_map
example_embedding_result = {
"object": "list",
@ -1743,30 +1744,39 @@ class TestPriceDataReloadAPI:
def test_reload_model_cost_map_admin_access(self, client_with_auth):
"""Test that admin users can access the reload endpoint"""
with patch(
"litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map"
) as mock_get_map:
mock_get_map.return_value = {
"gpt-3.5-turbo": {"input_cost_per_token": 0.001}
}
# Mock the database connection
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
# Save the original model_cost so the endpoint's direct assignment
# (litellm.model_cost = new_model_cost_map) does not contaminate
# subsequent tests running in the same worker process.
original_model_cost = litellm.model_cost.copy()
try:
with patch(
"litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map"
) as mock_get_map:
mock_get_map.return_value = {
"gpt-3.5-turbo": {"input_cost_per_token": 0.001}
}
# Mock the database connection
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
response = client_with_auth.post("/reload/model_cost_map")
response = client_with_auth.post("/reload/model_cost_map")
assert response.status_code == 200
data = response.json()
assert data["status"] == "success"
assert "message" in data
assert "timestamp" in data
assert "models_count" in data
# The new implementation immediately reloads and returns the count
assert (
"Price data reloaded successfully! 1 models updated."
in data["message"]
)
assert data["models_count"] == 1
assert response.status_code == 200
data = response.json()
assert data["status"] == "success"
assert "message" in data
assert "timestamp" in data
assert "models_count" in data
# The new implementation immediately reloads and returns the count
assert (
"Price data reloaded successfully! 1 models updated."
in data["message"]
)
assert data["models_count"] == 1
finally:
# Restore the full model cost map so subsequent tests are not affected
litellm.model_cost = original_model_cost
_invalidate_model_cost_lowercase_map()
def test_reload_model_cost_map_non_admin_access(self, client_with_auth):
"""Test that non-admin users cannot access the reload endpoint"""