diff --git a/tests/test_litellm/router_strategy/capability_router/test_capability_router.py b/tests/test_litellm/router_strategy/capability_router/test_capability_router.py index edc63b51e6e..00577184116 100644 --- a/tests/test_litellm/router_strategy/capability_router/test_capability_router.py +++ b/tests/test_litellm/router_strategy/capability_router/test_capability_router.py @@ -11,6 +11,7 @@ from litellm.router_strategy.capability_router.config import ( ) from litellm.router_strategy.capability_router.policy import select_capability_model from litellm.router_strategy.capability_router.prompts import build_classifier_response_schema +from litellm.types.router import Deployment, LiteLLM_Params def config() -> dict: @@ -121,6 +122,23 @@ def test_router_registers_capability_strategy() -> None: assert len(router.capability_routers["cost-router"]) == 1 +def test_router_explicitly_initializes_capability_strategy() -> None: + router = Router(model_list=[]) + deployment = Deployment( + model_name="cost-router", + litellm_params=LiteLLM_Params( + model="auto_router/capability_router", + capability_router_config=config(), + ), + ) + + assert router._is_capability_router_deployment(deployment.litellm_params) + + router.init_capability_router_deployment(deployment) + + assert len(router.capability_routers["cost-router"]) == 1 + + @pytest.mark.asyncio async def test_same_user_turn_reuses_cached_decision() -> None: router = Router(model_list=[])