diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py b/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py index 9846adf6142..11379ad2736 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py @@ -1,5 +1,6 @@ import os import sys +import asyncio from unittest.mock import AsyncMock, patch import pytest @@ -318,6 +319,36 @@ async def test_complexity_router_session_affinity_uses_router_configured_ttl(): assert any(call.kwargs.get("ttl") == 17 for call in cache.async_set_cache.call_args_list) +@pytest.mark.asyncio +async def test_complexity_router_session_affinity_expires_and_reselects(): + router = _complexity_router(session_affinity_ttl_seconds=1) + + with ( + patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new_callable=AsyncMock, + ) as mock_post, + patch( + "litellm.router_strategy.simple_shuffle.random.choice", + side_effect=_deterministic_choice(), + ), + ): + mock_post.return_value = _responses_mock() + first_response = await router.aresponses( + model="smart-router", + input="Hello", + metadata={"session_id": "expiry-session", "user_api_key_hash": "key-1"}, + ) + await asyncio.sleep(1.2) + second_response = await router.aresponses( + model="smart-router", + input="Follow-up", + metadata={"session_id": "expiry-session", "user_api_key_hash": "key-1"}, + ) + + assert second_response._hidden_params["model_id"] != first_response._hidden_params["model_id"] + + def test_complexity_router_registers_model_pool_groups_and_respects_disabled_affinity(): router = _complexity_router( session_affinity=True,