diff --git a/tests/router_unit_tests/test_pre_call_checks_optimization.py b/tests/router_unit_tests/test_pre_call_checks_optimization.py index bf70f06848f..1341d89accd 100644 --- a/tests/router_unit_tests/test_pre_call_checks_optimization.py +++ b/tests/router_unit_tests/test_pre_call_checks_optimization.py @@ -145,5 +145,37 @@ class TestPreCallChecksOptimization: ), "Second deployment ID changed!" +@pytest.mark.asyncio +async def test_async_wrapper_filters_like_sync(): + router = Router( + model_list=[ + { + "model_name": "test", + "litellm_params": {"model": "gpt-5-mini", "api_key": "sk-test"}, + "model_info": {"id": "small", "max_input_tokens": 50}, + }, + { + "model_name": "test", + "litellm_params": {"model": "gpt-5.5", "api_key": "sk-test"}, + "model_info": {"id": "large", "max_input_tokens": 10000}, + }, + ], + set_verbose=False, + enable_pre_call_checks=True, + ) + + deployments = router.get_model_list(model_name="test") + assert deployments is not None + + filtered = await router._pre_call_checks_async( + model="test", + healthy_deployments=deployments, + messages=[{"role": "user", "content": " ".join(["word"] * 100)}], + ) + + assert len(filtered) == 1 + assert filtered[0]["model_info"]["id"] == "large" + + if __name__ == "__main__": pytest.main([__file__, "-v"])