tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-11 00:25:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(tests): restore litellm.model_cost after TestPriceDataReloadIntegration tests

Browse Source 
test_complete_reload_flow and test_distributed_reload_check_function both
trigger code paths that assign a minimal stub dict to litellm.model_cost
(via the /reload/model_cost_map endpoint and _check_and_reload_model_cost_map).
Without restoring, subsequent tests in the same worker can't find gpt-4o
pricing and calculate spend=0.0 instead of the expected value.

Added try/finally save-and-restore of litellm.model_cost in both tests,
matching the pattern used in test_reload_model_cost_map_admin_access.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Julio Quinteros Pro
2026-02-19 01:42:48 -03:00
co-authored by Claude Sonnet 4.6
parent 155b79e53f
commit 68b83b1376
1 changed files with 38 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_litellm/proxy/test_proxy_server.py
+38 -28
View File
@@ -1999,22 +1999,27 @@ class TestPriceDataReloadIntegration:
"gpt-4": {"input_cost_per_token": 0.03, "output_cost_per_token": 0.06},
}
with patch(
"litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map"
) as mock_get_map:
mock_get_map.return_value = mock_cost_map
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 = mock_cost_map
# 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)
# 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)
# Test reload endpoint
response = client_with_auth.post("/reload/model_cost_map")
assert response.status_code == 200
# Test reload endpoint
response = client_with_auth.post("/reload/model_cost_map")
assert response.status_code == 200
# Test get endpoint
response = client_with_auth.get("/public/litellm_model_cost_map")
assert response.status_code == 200
# Test get endpoint
response = client_with_auth.get("/public/litellm_model_cost_map")
assert response.status_code == 200
finally:
litellm.model_cost = original_model_cost
_invalidate_model_cost_lowercase_map()
def test_distributed_reload_check_function(self):
"""Test the _check_and_reload_model_cost_map function"""
@@ -2054,23 +2059,28 @@ class TestPriceDataReloadIntegration:
mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config)
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
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}
}
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}
}
# Should reload due to force flag
asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma))
# Should reload due to force flag
asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma))
# Verify force_reload was reset to False
mock_prisma.db.litellm_config.upsert.assert_called()
call_args = mock_prisma.db.litellm_config.upsert.call_args
# The param_value is now a JSON string, so we need to parse it
param_value_json = call_args[1]["data"]["update"]["param_value"]
param_value_dict = json.loads(param_value_json)
assert param_value_dict["force_reload"] == False
# Verify force_reload was reset to False
mock_prisma.db.litellm_config.upsert.assert_called()
call_args = mock_prisma.db.litellm_config.upsert.call_args
# The param_value is now a JSON string, so we need to parse it
param_value_json = call_args[1]["data"]["update"]["param_value"]
param_value_dict = json.loads(param_value_json)
assert param_value_dict["force_reload"] == False
finally:
litellm.model_cost = original_model_cost
_invalidate_model_cost_lowercase_map()
def test_config_file_parsing(self):
"""Test parsing of config file with reload settings"""