mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-10 08:28:25 +00:00
ead221995e
Code outside the router now reaches it only through the entry points (registry, settings, health, select, apply, observe, register_plugin, retire) or the endpoint methods; the outside-member allowlist shrinks from 139 to 45. litellm/router.py no longer imports litellm.proxy: it asks the installed RouterHost (litellm/router_host.py) for vertex pass-through credentials, member auto-router authorization and auto-router prompt compression, and the proxy installs its host on import. Call sites where a model name can be missing keep the old answer instead of reaching an entry point with None.
774 lines
25 KiB
Python
774 lines
25 KiB
Python
"""Behavior pins for proxy_server background health-check helpers.
|
|
|
|
Pins covered:
|
|
- ``_get_process_rss_mb``
|
|
- ``_rss_mb_for_log``
|
|
- ``_run_direct_health_check_with_instrumentation``
|
|
- ``_schedule_background_health_check_db_save``
|
|
- ``_get_endpoint_exception_status``
|
|
- ``_write_health_state_to_router_cache``
|
|
- ``_adaptive_router_flusher_loop``
|
|
- ``_run_background_health_check``
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
|
|
import litellm.proxy.proxy_server as proxy_server
|
|
from litellm.constants import BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME
|
|
from litellm.router_api import RouterSettings
|
|
from litellm.proxy.proxy_server import (
|
|
_adaptive_router_flusher_loop,
|
|
_get_endpoint_exception_status,
|
|
_get_process_rss_mb,
|
|
_run_background_health_check,
|
|
_run_direct_health_check_with_instrumentation,
|
|
_rss_mb_for_log,
|
|
_schedule_background_health_check_db_save,
|
|
_write_health_state_to_router_cache,
|
|
)
|
|
|
|
from .conftest import normalize
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _get_process_rss_mb
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_get_process_rss_mb_returns_positive_float():
|
|
value = _get_process_rss_mb()
|
|
assert value is not None
|
|
assert normalize(
|
|
{
|
|
"value_present": value is not None,
|
|
"value_type": type(value).__name__,
|
|
"positive": value > 0,
|
|
}
|
|
) == {
|
|
"value_present": True,
|
|
"value_type": "float",
|
|
"positive": True,
|
|
}
|
|
|
|
|
|
def test_get_process_rss_mb_returns_none_when_resource_raises(monkeypatch):
|
|
import resource
|
|
|
|
def _boom(*_args, **_kwargs):
|
|
raise OSError("nope")
|
|
|
|
monkeypatch.setattr(resource, "getrusage", _boom)
|
|
assert _get_process_rss_mb() is None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _rss_mb_for_log
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_rss_mb_for_log_formats_numeric_value(monkeypatch):
|
|
monkeypatch.setattr(proxy_server, "_get_process_rss_mb", lambda: 100.5)
|
|
result = _rss_mb_for_log()
|
|
assert normalize(
|
|
{
|
|
"format": result,
|
|
"is_string": isinstance(result, str),
|
|
"contains_mb": "100.50" in result,
|
|
}
|
|
) == {
|
|
"format": "100.50",
|
|
"is_string": True,
|
|
"contains_mb": True,
|
|
}
|
|
|
|
|
|
def test_rss_mb_for_log_unknown_when_rss_missing(monkeypatch):
|
|
monkeypatch.setattr(proxy_server, "_get_process_rss_mb", lambda: None)
|
|
assert _rss_mb_for_log() == "unknown"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _run_direct_health_check_with_instrumentation
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_direct_health_check_with_instrumentation_returns_results(
|
|
monkeypatch,
|
|
):
|
|
expected = (["healthy_ep"], ["unhealthy_ep"], {"m1": Exception("boom")})
|
|
|
|
async def _fake_perform(model_list, details, max_concurrency, **kwargs):
|
|
return expected
|
|
|
|
monkeypatch.setattr(proxy_server, "perform_health_check", _fake_perform)
|
|
monkeypatch.setattr(
|
|
proxy_server,
|
|
"health_check_filter_kwargs_from_general_settings",
|
|
lambda _gs: {},
|
|
)
|
|
|
|
healthy, unhealthy, exceptions = await _run_direct_health_check_with_instrumentation(
|
|
model_list=[{"model_name": "gpt-4"}],
|
|
details=False,
|
|
max_concurrency=1,
|
|
instrumentation_context={"source": "test"},
|
|
)
|
|
|
|
assert normalize(
|
|
{
|
|
"healthy": healthy,
|
|
"unhealthy": unhealthy,
|
|
"exception_keys": list(exceptions.keys()),
|
|
}
|
|
) == {
|
|
"healthy": ["healthy_ep"],
|
|
"unhealthy": ["unhealthy_ep"],
|
|
"exception_keys": ["m1"],
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_direct_health_check_raises_non_kwarg_typeerror(monkeypatch):
|
|
async def _boom(model_list, details, max_concurrency, **kwargs):
|
|
raise TypeError("totally unrelated")
|
|
|
|
monkeypatch.setattr(proxy_server, "perform_health_check", _boom)
|
|
monkeypatch.setattr(
|
|
proxy_server,
|
|
"health_check_filter_kwargs_from_general_settings",
|
|
lambda _gs: {},
|
|
)
|
|
|
|
with pytest.raises(TypeError):
|
|
await _run_direct_health_check_with_instrumentation(
|
|
model_list=[],
|
|
details=False,
|
|
max_concurrency=1,
|
|
instrumentation_context={},
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _schedule_background_health_check_db_save
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schedule_background_health_check_db_save_creates_task(monkeypatch):
|
|
captured = {}
|
|
|
|
async def _fake_save(
|
|
prisma_client,
|
|
model_list,
|
|
healthy,
|
|
unhealthy,
|
|
start_time,
|
|
checked_by,
|
|
):
|
|
captured["prisma_client"] = prisma_client
|
|
captured["model_list"] = model_list
|
|
captured["healthy"] = healthy
|
|
captured["unhealthy"] = unhealthy
|
|
captured["checked_by"] = checked_by
|
|
|
|
import litellm.proxy.health_endpoints._health_endpoints as he
|
|
|
|
monkeypatch.setattr(he, "save_background_health_checks_to_db", _fake_save)
|
|
|
|
prisma_client = MagicMock()
|
|
shared_manager = SimpleNamespace(pod_id="pod-xyz")
|
|
|
|
_schedule_background_health_check_db_save(
|
|
prisma_client=prisma_client,
|
|
shared_health_manager=shared_manager,
|
|
model_list=[{"model_name": "gpt-4"}],
|
|
healthy_endpoints=[{"model_id": "h1"}],
|
|
unhealthy_endpoints=[{"model_id": "u1"}],
|
|
)
|
|
|
|
await asyncio.sleep(0)
|
|
|
|
assert normalize(
|
|
{
|
|
"prisma_present": captured.get("prisma_client") is prisma_client,
|
|
"checked_by": captured.get("checked_by"),
|
|
"healthy": captured.get("healthy"),
|
|
"unhealthy": captured.get("unhealthy"),
|
|
}
|
|
) == {
|
|
"prisma_present": True,
|
|
"checked_by": "pod-xyz",
|
|
"healthy": [{"model_id": "h1"}],
|
|
"unhealthy": [{"model_id": "u1"}],
|
|
}
|
|
|
|
|
|
def test_schedule_background_health_check_db_save_noop_when_prisma_none():
|
|
_schedule_background_health_check_db_save(
|
|
prisma_client=None,
|
|
shared_health_manager=None,
|
|
model_list=[],
|
|
healthy_endpoints=[],
|
|
unhealthy_endpoints=[],
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schedule_background_health_check_db_save_invalid_no_event_loop_raises(
|
|
monkeypatch,
|
|
):
|
|
async def _fake_save(*_args, **_kwargs):
|
|
return None
|
|
|
|
import litellm.proxy.health_endpoints._health_endpoints as he
|
|
|
|
monkeypatch.setattr(he, "save_background_health_checks_to_db", _fake_save)
|
|
|
|
def _broken_create_task(_coro):
|
|
raise RuntimeError("no running event loop")
|
|
|
|
monkeypatch.setattr(asyncio, "create_task", _broken_create_task)
|
|
|
|
with pytest.raises(RuntimeError):
|
|
_schedule_background_health_check_db_save(
|
|
prisma_client=MagicMock(),
|
|
shared_health_manager=None,
|
|
model_list=[],
|
|
healthy_endpoints=[],
|
|
unhealthy_endpoints=[],
|
|
)
|
|
|
|
|
|
def _lock_manager(redis_cache, acquired):
|
|
manager = MagicMock()
|
|
manager.redis_cache = redis_cache
|
|
manager.acquire_lock = AsyncMock(return_value=acquired)
|
|
manager.release_lock = AsyncMock()
|
|
return manager
|
|
|
|
|
|
def _capture_saves(monkeypatch, persisted=True):
|
|
saves = []
|
|
|
|
async def _fake_save(*_args, **kwargs):
|
|
saves.append(kwargs)
|
|
return persisted
|
|
|
|
import litellm.proxy.health_endpoints._health_endpoints as he
|
|
|
|
monkeypatch.setattr(he, "save_background_health_checks_to_db", _fake_save)
|
|
return saves
|
|
|
|
|
|
def _cancel_during_save(monkeypatch):
|
|
async def _fake_save(*_args, **_kwargs):
|
|
raise asyncio.CancelledError()
|
|
|
|
import litellm.proxy.health_endpoints._health_endpoints as he
|
|
|
|
monkeypatch.setattr(he, "save_background_health_checks_to_db", _fake_save)
|
|
|
|
|
|
def _schedule_with(lock_manager):
|
|
_schedule_background_health_check_db_save(
|
|
prisma_client=MagicMock(),
|
|
shared_health_manager=None,
|
|
model_list=[],
|
|
healthy_endpoints=[],
|
|
unhealthy_endpoints=[],
|
|
pod_lock_manager=lock_manager,
|
|
lock_ttl=300,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schedule_background_health_check_db_save_skips_a_window_another_pod_persisted(monkeypatch):
|
|
saves = _capture_saves(monkeypatch)
|
|
lock_manager = _lock_manager(redis_cache=MagicMock(), acquired=False)
|
|
|
|
_schedule_with(lock_manager)
|
|
await asyncio.sleep(0)
|
|
|
|
assert saves == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schedule_background_health_check_db_save_holds_the_window_lock_for_the_whole_interval(monkeypatch):
|
|
"""The lock is the "saved this window" marker: never reentrant, TTL = interval, and never released."""
|
|
saves = _capture_saves(monkeypatch)
|
|
lock_manager = _lock_manager(redis_cache=MagicMock(), acquired=True)
|
|
|
|
_schedule_with(lock_manager)
|
|
await asyncio.sleep(0)
|
|
|
|
assert normalize(
|
|
{
|
|
"saves": len(saves),
|
|
"lock_request": lock_manager.acquire_lock.await_args.kwargs,
|
|
"released": lock_manager.release_lock.await_count,
|
|
}
|
|
) == {
|
|
"saves": 1,
|
|
"lock_request": {
|
|
"cronjob_id": BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME,
|
|
"ttl": 300,
|
|
"allow_reentrant": False,
|
|
},
|
|
"released": 0,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schedule_background_health_check_db_save_releases_the_window_lock_when_the_save_reports_failure(
|
|
monkeypatch,
|
|
):
|
|
"""A failed save must not burn the window: release the lock so another pod's cycle can retry."""
|
|
saves = _capture_saves(monkeypatch, persisted=False)
|
|
lock_manager = _lock_manager(redis_cache=MagicMock(), acquired=True)
|
|
|
|
_schedule_with(lock_manager)
|
|
await asyncio.sleep(0)
|
|
|
|
assert normalize(
|
|
{
|
|
"saves": len(saves),
|
|
"release_request": lock_manager.release_lock.await_args.kwargs,
|
|
"release_count": lock_manager.release_lock.await_count,
|
|
}
|
|
) == {
|
|
"saves": 1,
|
|
"release_request": {"cronjob_id": BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME},
|
|
"release_count": 1,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schedule_background_health_check_db_save_releases_the_window_lock_when_the_save_is_cancelled(
|
|
monkeypatch,
|
|
):
|
|
"""A pod shutting down mid-save releases the lock instead of holding it until the TTL."""
|
|
_cancel_during_save(monkeypatch)
|
|
lock_manager = _lock_manager(redis_cache=MagicMock(), acquired=True)
|
|
|
|
_schedule_with(lock_manager)
|
|
await asyncio.sleep(0)
|
|
|
|
assert (
|
|
lock_manager.release_lock.await_args.kwargs,
|
|
lock_manager.release_lock.await_count,
|
|
) == ({"cronjob_id": BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME}, 1)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schedule_background_health_check_db_save_runs_ungated_without_redis(monkeypatch):
|
|
saves = _capture_saves(monkeypatch, persisted=False)
|
|
lock_manager = _lock_manager(redis_cache=None, acquired=True)
|
|
|
|
_schedule_with(lock_manager)
|
|
await asyncio.sleep(0)
|
|
|
|
assert (len(saves), lock_manager.acquire_lock.await_count, lock_manager.release_lock.await_count) == (1, 0, 0)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _get_endpoint_exception_status
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_get_endpoint_exception_status_prefers_live_exception():
|
|
endpoint = {"model_id": "m1", "exception_status": 999}
|
|
exceptions = {"m1": SimpleNamespace(status_code=429)}
|
|
status = _get_endpoint_exception_status(endpoint, exceptions)
|
|
assert normalize(
|
|
{
|
|
"input_endpoint": endpoint,
|
|
"exceptions_keys": list(exceptions.keys()),
|
|
"status": status,
|
|
}
|
|
) == {
|
|
"input_endpoint": {"model_id": "m1", "exception_status": 999},
|
|
"exceptions_keys": ["m1"],
|
|
"status": 429,
|
|
}
|
|
|
|
|
|
def test_get_endpoint_exception_status_falls_back_to_stored_int():
|
|
endpoint = {"model_id": "m-missing", "exception_status": 503}
|
|
assert _get_endpoint_exception_status(endpoint, {}) == 503
|
|
|
|
|
|
def test_get_endpoint_exception_status_default_500_when_no_data():
|
|
assert _get_endpoint_exception_status({}, {}) == 500
|
|
|
|
|
|
def test_get_endpoint_exception_status_invalid_endpoint_type_raises():
|
|
with pytest.raises(AttributeError):
|
|
_get_endpoint_exception_status(None, {}) # type: ignore[arg-type]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _write_health_state_to_router_cache
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_write_health_state_to_router_cache_sets_states(monkeypatch):
|
|
fake_router = MagicMock()
|
|
fake_router.settings.return_value = RouterSettings(
|
|
enable_health_check_routing=True, health_check_ignore_transient_errors=False, cooldown_time=30
|
|
)
|
|
health_checks = fake_router.health.return_value.health_checks
|
|
|
|
monkeypatch.setattr(proxy_server, "llm_router", fake_router)
|
|
|
|
fake_states = {"m1": {"is_healthy": True}, "m2": {"is_healthy": False}}
|
|
|
|
import litellm.proxy.health_check as hc
|
|
|
|
monkeypatch.setattr(hc, "build_deployment_health_states", lambda **_kw: fake_states)
|
|
|
|
import litellm.router_utils.cooldown_handlers as cd
|
|
|
|
monkeypatch.setattr(cd, "set_cooldown_deployments", lambda **_kw: None)
|
|
|
|
import litellm.router_utils.router_callbacks.track_deployment_metrics as tdm
|
|
|
|
monkeypatch.setattr(
|
|
tdm,
|
|
"increment_deployment_failures_for_current_minute",
|
|
lambda **_kw: None,
|
|
)
|
|
|
|
healthy = [{"model_id": "m1"}]
|
|
unhealthy = [{"model_id": "m2"}]
|
|
exceptions = {"m2": SimpleNamespace(status_code=500)}
|
|
|
|
_write_health_state_to_router_cache(healthy, unhealthy, exceptions)
|
|
|
|
health_checks.set_deployment_health_states.assert_called_once_with(fake_states)
|
|
|
|
call_args = health_checks.set_deployment_health_states.call_args[0][0]
|
|
assert normalize(
|
|
{
|
|
"states_keys": sorted(call_args.keys()),
|
|
"m1_healthy": call_args["m1"]["is_healthy"],
|
|
"m2_healthy": call_args["m2"]["is_healthy"],
|
|
}
|
|
) == {
|
|
"states_keys": ["m1", "m2"],
|
|
"m1_healthy": True,
|
|
"m2_healthy": False,
|
|
}
|
|
|
|
|
|
def test_write_health_state_to_router_cache_noop_when_router_none(monkeypatch):
|
|
monkeypatch.setattr(proxy_server, "llm_router", None)
|
|
_write_health_state_to_router_cache([], [], {})
|
|
|
|
|
|
def test_write_health_state_to_router_cache_noop_when_nothing_opted_in(monkeypatch):
|
|
"""Neither health-check routing nor the listing filter: write nothing."""
|
|
fake_router = MagicMock()
|
|
fake_router.settings.return_value = RouterSettings(
|
|
enable_health_check_routing=False, health_check_ignore_transient_errors=False
|
|
)
|
|
|
|
monkeypatch.setattr(proxy_server, "llm_router", fake_router)
|
|
monkeypatch.setattr(proxy_server, "general_settings", {})
|
|
|
|
_write_health_state_to_router_cache([{"model_id": "m1"}], [{"model_id": "m2"}], {})
|
|
|
|
fake_router.health.return_value.health_checks.set_deployment_health_states.assert_not_called()
|
|
|
|
|
|
def test_write_health_state_to_router_cache_populates_for_listing_filter(monkeypatch):
|
|
"""`model_list_healthy_only` needs the health cache, but must not start
|
|
cooling deployments down: that stays behind enable_health_check_routing."""
|
|
fake_router = MagicMock()
|
|
fake_router.settings.return_value = RouterSettings(
|
|
enable_health_check_routing=False, health_check_ignore_transient_errors=False, cooldown_time=30
|
|
)
|
|
|
|
monkeypatch.setattr(proxy_server, "llm_router", fake_router)
|
|
monkeypatch.setattr(proxy_server, "general_settings", {"model_list_healthy_only": True})
|
|
|
|
fake_states = {"m1": {"is_healthy": True}, "m2": {"is_healthy": False}}
|
|
|
|
import litellm.proxy.health_check as hc
|
|
|
|
monkeypatch.setattr(hc, "build_deployment_health_states", lambda **_kw: fake_states)
|
|
|
|
cooldowns: list[str] = []
|
|
|
|
import litellm.router_utils.cooldown_handlers as cd
|
|
|
|
monkeypatch.setattr(
|
|
cd,
|
|
"set_cooldown_deployments",
|
|
lambda **kw: cooldowns.append(kw.get("deployment")),
|
|
)
|
|
|
|
failures: list[str] = []
|
|
|
|
import litellm.router_utils.router_callbacks.track_deployment_metrics as tdm
|
|
|
|
monkeypatch.setattr(
|
|
tdm,
|
|
"increment_deployment_failures_for_current_minute",
|
|
lambda **kw: failures.append(kw.get("deployment_id")),
|
|
)
|
|
|
|
_write_health_state_to_router_cache(
|
|
[{"model_id": "m1"}],
|
|
[{"model_id": "m2"}],
|
|
{"m2": SimpleNamespace(status_code=500)},
|
|
)
|
|
|
|
fake_router.health.return_value.health_checks.set_deployment_health_states.assert_called_once_with(fake_states)
|
|
assert cooldowns == []
|
|
assert failures == []
|
|
|
|
|
|
def test_write_health_state_to_router_cache_swallows_internal_failures(monkeypatch):
|
|
"""The function logs and swallows exceptions so a bad cache call never crashes the loop."""
|
|
fake_router = MagicMock()
|
|
fake_router.settings.return_value = RouterSettings(
|
|
enable_health_check_routing=True, health_check_ignore_transient_errors=False
|
|
)
|
|
fake_router.health.return_value.health_checks.set_deployment_health_states.side_effect = RuntimeError(
|
|
"cache exploded"
|
|
)
|
|
|
|
monkeypatch.setattr(proxy_server, "llm_router", fake_router)
|
|
|
|
import litellm.proxy.health_check as hc
|
|
|
|
monkeypatch.setattr(
|
|
hc,
|
|
"build_deployment_health_states",
|
|
lambda **_kw: {"m1": {"is_healthy": True}},
|
|
)
|
|
|
|
_write_health_state_to_router_cache([{"model_id": "m1"}], [], {})
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _adaptive_router_flusher_loop
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_adaptive_router_flusher_loop_flushes_each_router(monkeypatch):
|
|
fake_ar = MagicMock()
|
|
fake_ar.state_loaded = True
|
|
fake_ar.queue.flush_state_to_db = AsyncMock()
|
|
fake_ar.queue.flush_session_to_db = AsyncMock()
|
|
|
|
from litellm.types.router import TaggedPreRoutingStrategy
|
|
|
|
fake_router = MagicMock()
|
|
fake_router.registry.return_value.plugins.return_value = {
|
|
"alpha": [TaggedPreRoutingStrategy(tags=(), strategy=fake_ar)]
|
|
}
|
|
|
|
monkeypatch.setattr(proxy_server, "llm_router", fake_router)
|
|
monkeypatch.setattr(proxy_server, "prisma_client", MagicMock())
|
|
|
|
# asyncio.sleep is awaited at the top of every iteration; raise CancelledError
|
|
# on the SECOND call so the first iteration completes its flush work.
|
|
call_count = {"n": 0}
|
|
_real_sleep = asyncio.sleep
|
|
|
|
async def _short_sleep(_seconds):
|
|
call_count["n"] += 1
|
|
if call_count["n"] >= 2:
|
|
raise asyncio.CancelledError()
|
|
await _real_sleep(0)
|
|
|
|
monkeypatch.setattr(proxy_server.asyncio, "sleep", _short_sleep)
|
|
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await _adaptive_router_flusher_loop()
|
|
|
|
assert fake_ar.queue.flush_state_to_db.await_count == 1
|
|
assert fake_ar.queue.flush_session_to_db.await_count == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_adaptive_router_flusher_loop_times_out_when_sleep_real(monkeypatch):
|
|
"""Confirms the loop is infinite — wait_for must raise TimeoutError."""
|
|
empty_router = MagicMock()
|
|
empty_router.registry.return_value.plugins.return_value = {}
|
|
monkeypatch.setattr(proxy_server, "llm_router", empty_router)
|
|
monkeypatch.setattr(proxy_server, "prisma_client", None)
|
|
|
|
# Bind the real asyncio.sleep before the patch so the replacement does not
|
|
# recurse into itself.
|
|
_real_sleep = asyncio.sleep
|
|
|
|
async def _instant_sleep(_seconds):
|
|
await _real_sleep(0)
|
|
|
|
monkeypatch.setattr(proxy_server.asyncio, "sleep", _instant_sleep)
|
|
|
|
with pytest.raises(asyncio.TimeoutError):
|
|
await asyncio.wait_for(_adaptive_router_flusher_loop(), timeout=0.2)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _run_background_health_check
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_background_health_check_returns_immediately_when_interval_invalid(
|
|
monkeypatch,
|
|
):
|
|
monkeypatch.setattr(proxy_server, "health_check_interval", None)
|
|
|
|
result = await _run_background_health_check()
|
|
|
|
assert normalize(
|
|
{
|
|
"result_is_none": result is None,
|
|
"loop_active": proxy_server.background_health_check_loop_active,
|
|
"interval": proxy_server.health_check_interval,
|
|
}
|
|
) == {
|
|
"result_is_none": True,
|
|
"loop_active": False,
|
|
"interval": None,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_background_health_check_runs_one_cycle_then_cancels(monkeypatch):
|
|
monkeypatch.setattr(proxy_server, "health_check_interval", 60)
|
|
monkeypatch.setattr(proxy_server, "health_check_concurrency", 1)
|
|
monkeypatch.setattr(proxy_server, "health_check_details", True)
|
|
monkeypatch.setattr(proxy_server, "use_shared_health_check", False)
|
|
monkeypatch.setattr(proxy_server, "redis_usage_cache", None)
|
|
monkeypatch.setattr(proxy_server, "prisma_client", None)
|
|
monkeypatch.setattr(proxy_server, "background_health_check_loop_active", False)
|
|
monkeypatch.setattr(
|
|
proxy_server,
|
|
"llm_model_list",
|
|
[{"model_name": "gpt-4", "model_info": {}}],
|
|
)
|
|
monkeypatch.setattr(
|
|
proxy_server,
|
|
"health_check_results",
|
|
{"healthy_endpoints": [], "unhealthy_endpoints": []},
|
|
)
|
|
|
|
async def _fake_direct(*_a, **_kw):
|
|
return ([{"model_id": "h"}], [{"model_id": "u"}], {})
|
|
|
|
monkeypatch.setattr(
|
|
proxy_server,
|
|
"_run_direct_health_check_with_instrumentation",
|
|
_fake_direct,
|
|
)
|
|
monkeypatch.setattr(proxy_server, "_schedule_background_health_check_db_save", lambda *a, **kw: None)
|
|
monkeypatch.setattr(proxy_server, "_write_health_state_to_router_cache", lambda *a, **kw: None)
|
|
monkeypatch.setattr(
|
|
proxy_server,
|
|
"health_check_filter_kwargs_from_general_settings",
|
|
lambda _gs: {},
|
|
)
|
|
|
|
sleep_calls = {"n": 0}
|
|
|
|
async def _stop_sleep(_seconds):
|
|
sleep_calls["n"] += 1
|
|
raise asyncio.CancelledError()
|
|
|
|
monkeypatch.setattr(proxy_server.asyncio, "sleep", _stop_sleep)
|
|
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await _run_background_health_check()
|
|
|
|
assert normalize(
|
|
{
|
|
"healthy_count": proxy_server.health_check_results["healthy_count"],
|
|
"unhealthy_count": proxy_server.health_check_results["unhealthy_count"],
|
|
"sleep_invoked": sleep_calls["n"] >= 1,
|
|
}
|
|
) == {
|
|
"healthy_count": 1,
|
|
"unhealthy_count": 1,
|
|
"sleep_invoked": True,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_background_health_check_probes_only_listed_model_groups(monkeypatch):
|
|
monkeypatch.setattr(proxy_server, "health_check_interval", 60)
|
|
monkeypatch.setattr(proxy_server, "health_check_concurrency", 1)
|
|
monkeypatch.setattr(proxy_server, "health_check_details", True)
|
|
monkeypatch.setattr(proxy_server, "use_shared_health_check", False)
|
|
monkeypatch.setattr(proxy_server, "redis_usage_cache", None)
|
|
monkeypatch.setattr(proxy_server, "prisma_client", None)
|
|
monkeypatch.setattr(proxy_server, "background_health_check_loop_active", False)
|
|
monkeypatch.setattr(
|
|
proxy_server,
|
|
"llm_router",
|
|
SimpleNamespace(
|
|
settings=lambda: RouterSettings(background_health_check_model_groups=frozenset({"prod-openai"}))
|
|
),
|
|
)
|
|
monkeypatch.setattr(
|
|
proxy_server,
|
|
"llm_model_list",
|
|
[
|
|
{"model_name": "prod-openai", "model_info": {"id": "listed-1"}},
|
|
{"model_name": "prod-openai", "model_info": {"id": "listed-2"}},
|
|
{"model_name": "internal-claude", "model_info": {"id": "unlisted-1"}},
|
|
{
|
|
"model_name": "prod-openai",
|
|
"model_info": {
|
|
"id": "listed-disabled",
|
|
"disable_background_health_check": True,
|
|
},
|
|
},
|
|
],
|
|
)
|
|
monkeypatch.setattr(
|
|
proxy_server,
|
|
"health_check_results",
|
|
{"healthy_endpoints": [], "unhealthy_endpoints": []},
|
|
)
|
|
|
|
probed = {}
|
|
|
|
async def _fake_direct(model_list, *_a, **_kw):
|
|
probed["ids"] = [m["model_info"]["id"] for m in model_list]
|
|
return ([], [], {})
|
|
|
|
monkeypatch.setattr(
|
|
proxy_server,
|
|
"_run_direct_health_check_with_instrumentation",
|
|
_fake_direct,
|
|
)
|
|
monkeypatch.setattr(proxy_server, "_schedule_background_health_check_db_save", lambda *a, **kw: None)
|
|
monkeypatch.setattr(proxy_server, "_write_health_state_to_router_cache", lambda *a, **kw: None)
|
|
monkeypatch.setattr(
|
|
proxy_server,
|
|
"health_check_filter_kwargs_from_general_settings",
|
|
lambda _gs: {},
|
|
)
|
|
|
|
async def _stop_sleep(_seconds):
|
|
raise asyncio.CancelledError()
|
|
|
|
monkeypatch.setattr(proxy_server.asyncio, "sleep", _stop_sleep)
|
|
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await _run_background_health_check()
|
|
|
|
assert probed["ids"] == ["listed-1", "listed-2"]
|