From bac9a36ca1ecf82f854dd4667d681fe9a27faa20 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 6 Aug 2026 12:25:13 +0000 Subject: [PATCH] test(proxy): isolate _runtime_registered_model_cost in distributed reload test A price data reload replays runtime registrations on top of the freshly adopted catalog. When a prior test in the worker had persisted a 'gpt-3.5-turbo' entry, this test's exact-equality assertion on litellm.model_cost['gpt-3.5-turbo'] would fail because register_model merges the built-in schema over the sparse mocked entry, producing over 100 fields instead of the single input_cost_per_token. Snapshot _runtime_registered_model_cost, clear it during the test, and restore it in finally. This matches the pattern already used in test_scheduled_reload_replays_runtime_registrations and test_swap_in_model_cost_map_counts_the_fetched_catalog_only in the same class. Co-authored-by: Krrish Dholakia --- tests/test_litellm/proxy/test_proxy_server.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index efd2ccb3e53..11019de7d0d 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -4006,6 +4006,7 @@ class TestPriceDataReloadIntegration: A revision this pod has not applied takes effect here even one minute into a 6h interval; a missing row is a no-op """ + from litellm import utils as litellm_utils from litellm.proxy.proxy_server import ProxyConfig proxy_config = ProxyConfig() @@ -4034,6 +4035,12 @@ class TestPriceDataReloadIntegration: from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded original_model_cost = litellm.model_cost.copy() + # A reload replays runtime registrations, so any gpt-3.5-turbo entry a + # prior test in this worker persisted would be re-registered on top of + # the mocked catalog and defeat the exact-equality check below. Isolate + # the registry for the duration of this test. + original_registry = dict(litellm_utils._runtime_registered_model_cost) + litellm_utils._runtime_registered_model_cost.clear() try: with ( patch("litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", new_callable=AsyncMock) as mock_get_map, @@ -4053,6 +4060,8 @@ class TestPriceDataReloadIntegration: assert proxy_config.model_cost_map_applied_revision == 4 finally: litellm.model_cost = original_model_cost + litellm_utils._runtime_registered_model_cost.clear() + litellm_utils._runtime_registered_model_cost.update(original_registry) _invalidate_model_cost_lowercase_map() def test_distributed_reload_ignores_already_applied_request(self):