From 25c5f0d993dc87069d18ad8b0a9b1fe8c57ada31 Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 2 Sep 2026 10:05:50 +0000 Subject: [PATCH 1/5] test: deflake JWT tamper assertions and fuzzy picker widget driver Tamper tests rewrote the last two base64url characters of the signature, which on roughly 1 in 250 RS256 tokens (1 in 1000 HS256) only touched padding bits, so the decoded signature was unchanged and still verified. Corrupt the decoded signature bytes instead. The fuzzy picker driver sent keys after fixed sleeps, so a slow worker could receive the filter text before the widget had highlighted the match. Wait on the widget's highlighted choice instead. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_session_credentials.py | 10 +++- .../test_session_token.py | 16 +++--- .../proxy/client/cli/autoroute/test_wizard.py | 51 +++++++++++++------ 3 files changed, 53 insertions(+), 24 deletions(-) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py index 00ff06ea082..992b1e8632b 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py @@ -5,6 +5,7 @@ from datetime import datetime, timedelta, timezone import pytest from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa +from jwt.utils import base64url_decode, base64url_encode from pydantic import SecretStr from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( @@ -52,6 +53,12 @@ def _refresh_token() -> str: return minted.token.get_secret_value() +def _corrupt_signature(token: str) -> str: + unsigned, signature = token.rsplit(".", 1) + raw = base64url_decode(signature) + return f"{unsigned}.{base64url_encode(bytes((raw[0] ^ 0x01,)) + raw[1:]).decode()}" + + def test_kdf_is_deterministic_and_key_length_is_256_bit(): again = session_keys_from_master_key(MASTER_KEY) assert again.signing_key.get_secret_value() == KEYS.signing_key.get_secret_value() @@ -109,8 +116,7 @@ def test_resolve_fails_expired_token_closed_and_flags_expiry(): def test_resolve_fails_tampered_token_closed_without_expiry_flag(): token = _access_token() - tampered = token[:-2] + ("aa" if not token.endswith("aa") else "bb") - result = resolve_session_bearer(f"Bearer {tampered}", KEYS, NOW) + result = resolve_session_bearer(f"Bearer {_corrupt_signature(token)}", KEYS, NOW) assert isinstance(result, SessionBearerInvalid) assert result.expired is False diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py index 2a59e6c1baa..321d6d0a1d2 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py @@ -6,6 +6,7 @@ import jwt import pytest from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa +from jwt.utils import base64url_decode, base64url_encode from pydantic import SecretStr, ValidationError from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( @@ -66,6 +67,12 @@ def _mint_refresh() -> str: return minted.token.get_secret_value() +def _corrupt_signature(token: str) -> str: + unsigned, signature = token.rsplit(".", 1) + raw = base64url_decode(signature) + return f"{unsigned}.{base64url_encode(bytes((raw[0] ^ 0x01,)) + raw[1:]).decode()}" + + def _sign_claims(payload: dict, prefix: str = SESSION_TOKEN_PREFIX, keys: SessionKeys = KEYS) -> str: return prefix + jwt.encode(payload, keys.signing_key.get_secret_value(), algorithm="HS256") @@ -138,8 +145,7 @@ def test_still_valid_one_second_before_expiry(): def test_tampered_signature_is_bad_signature(): token = _mint_access() - tampered = token[:-2] + ("aa" if not token.endswith("aa") else "bb") - assert isinstance(open_session_token(tampered, KEYS, NOW), SessionBadSignature) + assert isinstance(open_session_token(_corrupt_signature(token), KEYS, NOW), SessionBadSignature) def test_key_rotation_invalidates_outstanding_tokens(): @@ -329,8 +335,7 @@ def test_rs256_tampered_signature_is_bad_signature(): minted = mint_session_token(PRINCIPAL, RSA_KEYS, NOW) assert isinstance(minted, MintedSessionToken) token = minted.token.get_secret_value() - tampered = token[:-2] + ("aa" if not token.endswith("aa") else "bb") - assert isinstance(open_session_token(tampered, RSA_KEYS, NOW), SessionBadSignature) + assert isinstance(open_session_token(_corrupt_signature(token), RSA_KEYS, NOW), SessionBadSignature) def test_rs256_expired_token_is_expired(): @@ -413,8 +418,7 @@ def test_rotation_window_still_enforces_expiry_and_tamper_on_the_previous_key(): ) after = NOW + timedelta(seconds=SESSION_TTL_SECONDS + 1) assert isinstance(open_session_token(token, rotated, after), SessionExpired) - tampered = token[:-2] + ("aa" if not token.endswith("aa") else "bb") - assert isinstance(open_session_token(tampered, rotated, NOW), SessionBadSignature) + assert isinstance(open_session_token(_corrupt_signature(token), rotated, NOW), SessionBadSignature) def test_weak_or_garbage_private_key_pem_rejected_at_construction(): diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_wizard.py b/tests/test_litellm/proxy/client/cli/autoroute/test_wizard.py index a17fed36f52..fc6de53cb9e 100644 --- a/tests/test_litellm/proxy/client/cli/autoroute/test_wizard.py +++ b/tests/test_litellm/proxy/client/cli/autoroute/test_wizard.py @@ -1,5 +1,5 @@ import asyncio -from typing import Any, Dict, List, Tuple +from typing import Any, Dict, List, Optional, Tuple from unittest.mock import patch import click @@ -7,7 +7,8 @@ import pytest import yaml from click.testing import CliRunner from InquirerPy.base.control import Choice -from prompt_toolkit.application import create_app_session +from InquirerPy.prompts.fuzzy import InquirerPyFuzzyControl +from prompt_toolkit.application import AppSession, create_app_session from prompt_toolkit.input import create_pipe_input from prompt_toolkit.output import DummyOutput @@ -283,27 +284,45 @@ class TestRunConfigureWizardNotInteractive: assert not config_path.exists() +def _highlighted_choice(session: AppSession) -> Optional[str]: + if session.app is None: + return None + controls = [c for c in session.app.layout.find_all_controls() if isinstance(c, InquirerPyFuzzyControl)] + if not controls or controls[0].choice_count == 0: + return None + return controls[0].selection["name"] + + +async def _wait_until_highlighted(session: AppSession, name: str) -> None: + async def _poll() -> None: + while _highlighted_choice(session) != name: + await asyncio.sleep(0.01) + + await asyncio.wait_for(_poll(), timeout=5) + + def _drive_fuzzy_pick( models: Tuple[DiscoveredModel, ...], prompt_label: str, multiselect: bool, - key_events: List[Tuple[str, float]], + key_events: List[Tuple[str, Optional[str]]], ) -> List[str]: """Drives the real InquirerPy fuzzy prompt through prompt_toolkit's own test input/output, exercising the actual widget (filtering, tab-to-toggle, enter-to-confirm) rather than mocking it away. asyncio.to_thread propagates the create_app_session context into the worker thread - running _fuzzy_pick's synchronous .execute() call.""" + running _fuzzy_pick's synchronous .execute() call. Each key event names the choice the widget + must highlight before the next key is sent (None sends the next key immediately).""" async def _run() -> List[str]: with create_pipe_input() as pipe_input: - with create_app_session(input=pipe_input, output=DummyOutput()): + with create_app_session(input=pipe_input, output=DummyOutput()) as session: task = asyncio.ensure_future( asyncio.to_thread(wizard_module._fuzzy_pick, models, prompt_label, multiselect) ) - await asyncio.sleep(0.05) - for text, delay in key_events: + for text, highlighted in key_events: pipe_input.send_text(text) - await asyncio.sleep(delay) + if highlighted is not None: + await _wait_until_highlighted(session, highlighted) return await task return asyncio.run(_run()) @@ -315,13 +334,13 @@ class TestFuzzyPickWidget: def test_single_select_filters_and_returns_highlighted_match(self): result = _drive_fuzzy_pick( - self._models(), "test", multiselect=False, key_events=[("model-13", 0.3), ("\r", 0.1)] + self._models(), "test", multiselect=False, key_events=[("model-13", "model-13"), ("\r", None)] ) assert result == ["model-13"] def test_multiselect_requires_tab_to_toggle_before_enter(self): result = _drive_fuzzy_pick( - self._models(), "test", multiselect=True, key_events=[("model-7", 0.3), ("\t", 0.1), ("\r", 0.1)] + self._models(), "test", multiselect=True, key_events=[("model-7", "model-7"), ("\t", None), ("\r", None)] ) assert result == ["model-7"] @@ -331,12 +350,12 @@ class TestFuzzyPickWidget: "test", multiselect=True, key_events=[ - ("model-3", 0.3), - ("\t", 0.1), - *[("\x7f", 0.02) for _ in range("model-3".__len__())], - ("model-15", 0.3), - ("\t", 0.1), - ("\r", 0.1), + ("model-3", "model-3"), + ("\t", None), + ("\x7f" * len("model-3"), None), + ("model-15", "model-15"), + ("\t", None), + ("\r", None), ], ) assert set(result) == {"model-3", "model-15"} From 753bea360e8b0921b9a5fe02ba860f558a147e39 Mon Sep 17 00:00:00 2001 From: mateo Date: Thu, 3 Sep 2026 10:38:20 +0000 Subject: [PATCH 2/5] test: deflake guardrail mapping leak, tag routing randomness, and liveliness timing TestStreamingScanDedup restored the reduced module-level translation mapping on teardown via monkeypatch, so under --dist=loadscope the worker that ran only that class carried the reduced mapping into the streaming block test modules. Tag routing tests now assert the eligible deployment set directly instead of sampling ten random picks. The liveliness latency check measures steady-state polls after a warm-up request rather than the first request through a fresh app. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_unified_guardrail.py | 10 +++--- .../health_endpoints/test_health_endpoints.py | 30 +++++++--------- .../test_router_tag_routing.py | 36 +++++++++---------- 3 files changed, 33 insertions(+), 43 deletions(-) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index a28a2a71613..ceaf7c49595 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -2037,12 +2037,10 @@ class TestStreamingScanDedup: guardrail already cleared. Regression for LIT-6692.""" @pytest.fixture(autouse=True) - def _use_real_mappings(self, monkeypatch): - monkeypatch.setattr( - unified_module, - "endpoint_guardrail_translation_mappings", - load_guardrail_translation_mappings(), - ) + def _use_real_mappings(self): + unified_module.endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings() + yield + unified_module.endpoint_guardrail_translation_mappings = None @pytest.mark.asyncio async def test_chat_terminal_chunk_on_sampled_index_is_scanned_once(self): diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index e3f71692c78..619f6736ea3 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -1,6 +1,7 @@ import asyncio import json import time +from typing import Final from datetime import datetime, timedelta from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch @@ -1189,27 +1190,22 @@ def test_health_liveliness_endpoint(proxy_client): Test that /health/liveliness endpoint returns 200 OK with "I'm alive!" message. This is a critical orchestration endpoint that must be simple and fast. """ - # Measure the time taken for the health check call - start_time = time.perf_counter() + warm_up: Final = proxy_client.get("/health/liveliness") + assert warm_up.status_code == 200, f"Expected 200 OK, got {warm_up.status_code}: {warm_up.text}" - # Make GET request to /health/liveliness - response = proxy_client.get("/health/liveliness") + def _timed_poll() -> tuple[float, httpx.Response]: + start_time: Final = time.perf_counter() + response: Final = proxy_client.get("/health/liveliness") + return (time.perf_counter() - start_time) * 1000, response - end_time = time.perf_counter() - duration_ms = (end_time - start_time) * 1000 + polls: Final = tuple(_timed_poll() for _ in range(5)) - # Assert response status - assert response.status_code == 200, f"Expected 200 OK, got {response.status_code}: {response.text}" + for _, response in polls: + assert response.status_code == 200, f"Expected 200 OK, got {response.status_code}: {response.text}" + assert response.json() == "I'm alive!", f"Expected 'I'm alive!' message, got: {response.json()}" - # Assert response content (FastAPI JSON-encodes the string) - assert response.json() == "I'm alive!", f"Expected 'I'm alive!' message, got: {response.json()}" - - # Verify response is fast (should be < 100ms for a simple endpoint) - # This is critical for orchestration systems that poll frequently - assert duration_ms < 100, f"Health check took {duration_ms:.2f}ms, expected < 100ms for a simple endpoint" - - # Log the duration for visibility (useful for CI/CD monitoring) - print(f"\n/health/liveliness response time: {duration_ms:.2f}ms") + fastest_ms: Final = min(duration_ms for duration_ms, _ in polls) + assert fastest_ms < 100, f"Fastest of {len(polls)} health checks took {fastest_ms:.2f}ms, expected < 100ms" def test_health_liveness_endpoint(proxy_client): diff --git a/tests/test_litellm/router_strategy/test_router_tag_routing.py b/tests/test_litellm/router_strategy/test_router_tag_routing.py index b33bd912be9..27f871ed39f 100644 --- a/tests/test_litellm/router_strategy/test_router_tag_routing.py +++ b/tests/test_litellm/router_strategy/test_router_tag_routing.py @@ -5,10 +5,22 @@ import pytest import logging +from typing import Final import litellm from litellm._logging import verbose_logger +from litellm.router_strategy.tag_based_routing import get_deployments_for_tag + + +async def _eligible_deployment_ids(router: litellm.Router, model: str, tags: list[str]) -> set[str]: + eligible: Final = await get_deployments_for_tag( + llm_router_instance=router, + model=model, + healthy_deployments=router.get_model_list(model_name=model) or [], + request_kwargs={"metadata": {"tags": tags}}, + ) + return {deployment["model_info"]["id"] for deployment in eligible} @pytest.mark.asyncio() @@ -850,17 +862,9 @@ async def test_negation_regex_pattern_treated_as_literal(): # The regex-like string matches no deployment tag literally, so all # candidates survive and both model IDs are reachable. - seen_ids = set() - for _ in range(10): - response = await router.acompletion( - model="gpt-4", - messages=[{"role": "user", "content": "hi"}], - metadata={"tags": ["!provider:(anthropic|openai)"]}, - mock_response="hi", - ) - seen_ids.add(response._hidden_params["model_id"]) + eligible_ids: Final = await _eligible_deployment_ids(router, "gpt-4", ["!provider:(anthropic|openai)"]) - assert seen_ids == {"anthropic-model", "openai-model"} + assert eligible_ids == {"anthropic-model", "openai-model"} @pytest.mark.asyncio() @@ -1281,17 +1285,9 @@ async def test_chain_enable_tag_filtering_false_overrides_router_level_true(): enable_tag_filtering=True, ) - seen_ids = set() - for _ in range(10): - response = await router.acompletion( - model="gpt-4", - messages=[{"role": "user", "content": "hi"}], - metadata={"tags": ["teamA"]}, - mock_response="hi", - ) - seen_ids.add(response._hidden_params["model_id"]) + eligible_ids: Final = await _eligible_deployment_ids(router, "gpt-4", ["teamA"]) - assert seen_ids == {"team-a-deployment", "team-b-deployment"} + assert eligible_ids == {"team-a-deployment", "team-b-deployment"} @pytest.mark.asyncio() From 03a82823fc46e19fde89021acbd9095edb94fa79 Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 4 Sep 2026 10:07:34 +0000 Subject: [PATCH 3/5] test: deflake redis loop-stall burst test and pre-commit interrupt cleanup The redis breaker test raced the event loop: the fake call had to still be pending when a real time.sleep stall began, which needs the loop to get from scheduling to the stall in under 1ms. The fake now holds its answer behind an asyncio.Event so the whole burst times out deterministically. The pre-commit interrupt test found a real leak: lint_dashboard creates its eslint report with mktemp and only removed it on the happy path, so an interrupt landing during the whole-folder eslint run left the file behind. The subshell now removes it from an EXIT trap, and the test drives the interrupt while that eslint run is in flight. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- scripts/pre_commit_lint.sh | 3 ++- tests/test_litellm/caching/test_redis_cache.py | 17 +++++++---------- tests/test_litellm/test_pre_commit_lint.py | 9 ++++++++- 3 files changed, 17 insertions(+), 12 deletions(-) diff --git a/scripts/pre_commit_lint.sh b/scripts/pre_commit_lint.sh index ff553be6461..d38a0eee3de 100755 --- a/scripts/pre_commit_lint.sh +++ b/scripts/pre_commit_lint.sh @@ -142,6 +142,8 @@ fi lint_dashboard() { ( + trap 'exit 143' TERM + trap 'rm -f "${report:-}"' EXIT rc=0 prettier_rel=() eslint_rel=() @@ -168,7 +170,6 @@ EOF report=$(mktemp) npx eslint . -f json -o "$report" || true node scripts/check-lint-budgets.mjs "$report" eslint-budgets.json || rc=1 - rm -f "$report" exit $rc ) } diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index e4724ff8705..dd87bf0cf0f 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -822,30 +822,27 @@ async def test_event_loop_stall_timeout_burst_keeps_breaker_closed(): Every operation already waiting on the loop times out together when the loop resumes, so a purely consecutive threshold is satisfied instantly even though the Redis on the - other end (here an in-process fake that answers immediately) is healthy. + other end is healthy. The stall is modelled by holding the fake Redis's answers back + until the whole burst has hit its client timeout, then releasing them. """ - import time as time_mod - from litellm.caching.redis_cache import RedisCircuitBreaker, _run_under_circuit_breaker breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, timeout_min_duration=5.0) + loop_resumed = asyncio.Event() async def healthy_redis_call_with_client_timeout(): - return await asyncio.wait_for(asyncio.sleep(0.001, result="ok"), timeout=0.05) - - async def stall_the_loop(): - await asyncio.sleep(0) - time_mod.sleep(0.2) + await asyncio.wait_for(loop_resumed.wait(), timeout=0.05) + return "ok" results = await asyncio.gather( *(_run_under_circuit_breaker(breaker, "op", healthy_redis_call_with_client_timeout) for _ in range(8)), - stall_the_loop(), return_exceptions=True, ) timeouts = [r for r in results if isinstance(r, asyncio.TimeoutError)] - assert len(timeouts) >= breaker.failure_threshold, "the stall must time out a full burst" + assert len(timeouts) == 8, "the stall must time out the whole burst" assert breaker.is_open() is False, "a healthy Redis behind one loop stall must stay in the pool" + loop_resumed.set() assert await _run_under_circuit_breaker(breaker, "op", healthy_redis_call_with_client_timeout) == "ok" diff --git a/tests/test_litellm/test_pre_commit_lint.py b/tests/test_litellm/test_pre_commit_lint.py index 5ea0e79a196..aa2260e89ea 100644 --- a/tests/test_litellm/test_pre_commit_lint.py +++ b/tests/test_litellm/test_pre_commit_lint.py @@ -53,6 +53,12 @@ case "$*" in "eslint --no-warn-ignored"*) [ "${STUB_FAIL:-}" = "eslint" ] && exit 1 ;; + "eslint . -f json"*) + if [ -n "${STUB_HANG_DIR:-}" ]; then + touch "$STUB_HANG_DIR/eslint_report.started" + sleep 60 + fi + ;; esac exit 0 """ @@ -340,11 +346,12 @@ def test_interrupt_kills_background_jobs_and_removes_logs(tmp_path: Path) -> Non ) try: assert _wait_until((hang_dir / "make.started").exists, 10) + assert _wait_until((hang_dir / "eslint_report.started").exists, 10) os.killpg(proc.pid, signal.SIGINT) assert proc.wait(timeout=10) != 0 make_pid = int((hang_dir / "make.pid").read_text()) assert _wait_until(lambda: _pid_gone(make_pid), 5) - assert list(tmp_dir.iterdir()) == [] + assert _wait_until(lambda: not any(tmp_dir.iterdir()), 5), list(tmp_dir.iterdir()) finally: with suppress(ProcessLookupError, PermissionError): os.killpg(proc.pid, signal.SIGTERM) From 5bd4da0389f42d1b0f32f27c2061d17523f7cdab Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 19:16:51 -0700 Subject: [PATCH 4/5] test(health): score the liveliness probe on the median of five warm polls --- .../proxy/health_endpoints/test_health_endpoints.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index c9cd4e0d9a1..e9e58347337 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -1205,8 +1205,9 @@ def test_health_liveliness_endpoint(proxy_client): assert response.status_code == 200, f"Expected 200 OK, got {response.status_code}: {response.text}" assert response.json() == "I'm alive!", f"Expected 'I'm alive!' message, got: {response.json()}" - fastest_ms: Final = min(duration_ms for duration_ms, _ in polls) - assert fastest_ms < 100, f"Fastest of {len(polls)} health checks took {fastest_ms:.2f}ms, expected < 100ms" + durations_ms: Final = tuple(sorted(duration_ms for duration_ms, _ in polls)) + median_ms: Final = durations_ms[len(durations_ms) // 2] + assert median_ms < 100, f"Median of {len(polls)} health checks took {median_ms:.2f}ms, expected < 100ms" def test_health_liveness_endpoint(proxy_client): From 41c8969f0aee5769f5feded4bc5e7ff8723db469 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 19:54:40 -0700 Subject: [PATCH 5/5] test(router): drive tag routing tests through acompletion until both deployments are seen --- .../test_router_tag_routing.py | 30 ++++++++++--------- 1 file changed, 16 insertions(+), 14 deletions(-) diff --git a/tests/test_litellm/router_strategy/test_router_tag_routing.py b/tests/test_litellm/router_strategy/test_router_tag_routing.py index 27f871ed39f..59a59c7e16d 100644 --- a/tests/test_litellm/router_strategy/test_router_tag_routing.py +++ b/tests/test_litellm/router_strategy/test_router_tag_routing.py @@ -2,25 +2,25 @@ # This tests litellm router -import pytest - import logging from typing import Final +import pytest import litellm from litellm._logging import verbose_logger -from litellm.router_strategy.tag_based_routing import get_deployments_for_tag -async def _eligible_deployment_ids(router: litellm.Router, model: str, tags: list[str]) -> set[str]: - eligible: Final = await get_deployments_for_tag( - llm_router_instance=router, - model=model, - healthy_deployments=router.get_model_list(model_name=model) or [], - request_kwargs={"metadata": {"tags": tags}}, +async def _routed_model_ids( + router: litellm.Router, tags: list[str], remaining: frozenset[str], attempts: int = 100 +) -> frozenset[str]: + if not remaining or attempts == 0: + return frozenset() + response: Final = await router.acompletion( + model="gpt-4", messages=[{"role": "user", "content": "hi"}], metadata={"tags": tags}, mock_response="hi" ) - return {deployment["model_info"]["id"] for deployment in eligible} + seen: Final = frozenset({response._hidden_params["model_id"]}) + return seen | await _routed_model_ids(router, tags, remaining - seen, attempts - 1) @pytest.mark.asyncio() @@ -862,9 +862,10 @@ async def test_negation_regex_pattern_treated_as_literal(): # The regex-like string matches no deployment tag literally, so all # candidates survive and both model IDs are reachable. - eligible_ids: Final = await _eligible_deployment_ids(router, "gpt-4", ["!provider:(anthropic|openai)"]) + expected: Final = frozenset({"anthropic-model", "openai-model"}) + routed_ids: Final = await _routed_model_ids(router, ["!provider:(anthropic|openai)"], expected) - assert eligible_ids == {"anthropic-model", "openai-model"} + assert routed_ids == expected @pytest.mark.asyncio() @@ -1285,9 +1286,10 @@ async def test_chain_enable_tag_filtering_false_overrides_router_level_true(): enable_tag_filtering=True, ) - eligible_ids: Final = await _eligible_deployment_ids(router, "gpt-4", ["teamA"]) + expected: Final = frozenset({"team-a-deployment", "team-b-deployment"}) + routed_ids: Final = await _routed_model_ids(router, ["teamA"], expected) - assert eligible_ids == {"team-a-deployment", "team-b-deployment"} + assert routed_ids == expected @pytest.mark.asyncio()