From 509b1c48f143b77bb74fb118879504f765b5bf43 Mon Sep 17 00:00:00 2001 From: Arya Maddel <84805906+aryamaddel@users.noreply.github.com> Date: Tue, 5 May 2026 21:16:41 +0530 Subject: [PATCH] test: add tests for model_info=None in router_strategy callbacks --- .../local_testing/test_least_busy_routing.py | 58 +++++++++++++++++++ .../local_testing/test_lowest_cost_routing.py | 56 ++++++++++++++++++ .../local_testing/test_tpm_rpm_routing_v2.py | 17 ++++++ 3 files changed, 131 insertions(+) diff --git a/tests/local_testing/test_least_busy_routing.py b/tests/local_testing/test_least_busy_routing.py index 0f4f6923a19..94722eecb0f 100644 --- a/tests/local_testing/test_least_busy_routing.py +++ b/tests/local_testing/test_least_busy_routing.py @@ -26,6 +26,64 @@ from litellm.router_strategy.least_busy import LeastBusyLoggingHandler ### UNIT TESTS FOR LEAST BUSY LOGGING ### +def test_least_busy_sync_success_model_info_none(): + """log_success_event should not crash when model_info is None.""" + logger = LeastBusyLoggingHandler(router_cache=DualCache()) + kwargs = { + "litellm_params": { + "metadata": {"model_group": "gpt-4"}, + "model_info": None, + } + } + logger.log_success_event( + kwargs=kwargs, response_obj={}, start_time=time.time(), end_time=time.time() + ) + + +def test_least_busy_sync_failure_model_info_none(): + """log_failure_event should not crash when model_info is None.""" + logger = LeastBusyLoggingHandler(router_cache=DualCache()) + kwargs = { + "litellm_params": { + "metadata": {"model_group": "gpt-4"}, + "model_info": None, + } + } + logger.log_failure_event( + kwargs=kwargs, response_obj={}, start_time=time.time(), end_time=time.time() + ) + + +@pytest.mark.asyncio +async def test_least_busy_async_success_model_info_none(): + """async_log_success_event should not crash when model_info is None.""" + logger = LeastBusyLoggingHandler(router_cache=DualCache()) + kwargs = { + "litellm_params": { + "metadata": {"model_group": "gpt-4"}, + "model_info": None, + } + } + await logger.async_log_success_event( + kwargs=kwargs, response_obj={}, start_time=time.time(), end_time=time.time() + ) + + +@pytest.mark.asyncio +async def test_least_busy_async_failure_model_info_none(): + """async_log_failure_event should not crash when model_info is None.""" + logger = LeastBusyLoggingHandler(router_cache=DualCache()) + kwargs = { + "litellm_params": { + "metadata": {"model_group": "gpt-4"}, + "model_info": None, + } + } + await logger.async_log_failure_event( + kwargs=kwargs, response_obj={}, start_time=time.time(), end_time=time.time() + ) + + def test_model_added(): test_cache = DualCache() least_busy_logger = LeastBusyLoggingHandler(router_cache=test_cache) diff --git a/tests/local_testing/test_lowest_cost_routing.py b/tests/local_testing/test_lowest_cost_routing.py index 4e8b06fb628..573defefcae 100644 --- a/tests/local_testing/test_lowest_cost_routing.py +++ b/tests/local_testing/test_lowest_cost_routing.py @@ -20,6 +20,62 @@ from litellm.caching.caching import DualCache ### UNIT TESTS FOR cost ROUTING ### +def test_lowest_cost_sync_success_model_info_none(): + """log_success_event should not crash when model_info is None.""" + logger = LowestCostLoggingHandler(router_cache=DualCache()) + kwargs = { + "litellm_params": { + "metadata": {"model_group": "gpt-4"}, + "model_info": None, + } + } + logger.log_success_event( + kwargs=kwargs, response_obj={}, start_time=time.time(), end_time=time.time() + ) + + +def test_lowest_cost_sync_failure_model_info_none(): + """log_failure_event should not crash when model_info is None.""" + logger = LowestCostLoggingHandler(router_cache=DualCache()) + kwargs = { + "litellm_params": { + "metadata": {"model_group": "gpt-4"}, + "model_info": None, + } + } + logger.log_failure_event( + kwargs=kwargs, response_obj={}, start_time=time.time(), end_time=time.time() + ) + + +def test_lowest_cost_sync_success_model_info_missing(): + """log_success_event should not crash when model_info key is missing.""" + logger = LowestCostLoggingHandler(router_cache=DualCache()) + kwargs = { + "litellm_params": { + "metadata": {"model_group": "gpt-4"}, + } + } + logger.log_success_event( + kwargs=kwargs, response_obj={}, start_time=time.time(), end_time=time.time() + ) + + +@pytest.mark.asyncio +async def test_lowest_cost_async_success_model_info_none(): + """async_log_success_event should not crash when model_info is None.""" + logger = LowestCostLoggingHandler(router_cache=DualCache()) + kwargs = { + "litellm_params": { + "metadata": {"model_group": "gpt-4"}, + "model_info": None, + } + } + await logger.async_log_success_event( + kwargs=kwargs, response_obj={}, start_time=time.time(), end_time=time.time() + ) + + @pytest.mark.asyncio async def test_get_available_deployments(): test_cache = DualCache() diff --git a/tests/local_testing/test_tpm_rpm_routing_v2.py b/tests/local_testing/test_tpm_rpm_routing_v2.py index 211af566424..0569773cae0 100644 --- a/tests/local_testing/test_tpm_rpm_routing_v2.py +++ b/tests/local_testing/test_tpm_rpm_routing_v2.py @@ -27,6 +27,9 @@ from litellm.caching.caching import DualCache from litellm.router_strategy.lowest_tpm_rpm_v2 import ( LowestTPMLoggingHandler_v2 as LowestTPMLoggingHandler, ) +from litellm.router_strategy.lowest_tpm_rpm import ( + LowestTPMLoggingHandler as LowestTPMLoggingHandler_v1, +) from litellm.utils import get_utc_datetime from create_mock_standard_logging_payload import create_standard_logging_payload @@ -37,6 +40,20 @@ from create_mock_standard_logging_payload import create_standard_logging_payload """ +def test_lowest_tpm_rpm_v1_model_info_none(): + """log_success_event should not crash when model_info is None (v1 handler).""" + logger = LowestTPMLoggingHandler_v1(router_cache=DualCache()) + kwargs = { + "litellm_params": { + "metadata": {"model_group": "gpt-4"}, + "model_info": None, + } + } + logger.log_success_event( + kwargs=kwargs, response_obj={}, start_time=time.time(), end_time=time.time() + ) + + def test_tpm_rpm_updated(): test_cache = DualCache() lowest_tpm_logger = LowestTPMLoggingHandler(router_cache=test_cache)