litellm/tests/integration/_support/runtime.py
yuneng-jiang b67003162a
test(integration): wait for every worker to report the websocket session cap instead of sleeping (#47412)
The override test slept a fixed 10 s after writing responses_websocket_session_limit_seconds and then
opened sockets on both workers. A worker whose config resync ran later than that kept the old 7200 s
cap, so its socket never closed and the test failed after 75 s

A test-only /health/integration-general-setting route returns the value a worker holds along with its
pid. The test polls it over fresh connections until every worker reports the new cap, and after the
delete until every worker reports it gone
2026-10-19 15:35:29 -07:00

49 lines
1.8 KiB
Python

import asyncio
import os
from concurrent.futures import ThreadPoolExecutor
from typing import Final
from fastapi import Response
from fastapi.responses import JSONResponse
from pydantic import TypeAdapter
READINESS_BODY: Final = TypeAdapter(dict[str, object])
def configure_executor() -> None:
asyncio.get_running_loop().set_default_executor(
ThreadPoolExecutor(max_workers=32, thread_name_prefix="integration-provider")
)
async def worker_readiness() -> JSONResponse:
from litellm.proxy.health_endpoints._health_endpoints import health_readiness
response: Final = Response()
body: Final = READINESS_BODY.validate_python(await health_readiness(response))
return JSONResponse(
content=body, status_code=response.status_code, headers={"x-integration-worker-pid": str(os.getpid())}
)
async def worker_deployment_ids(model_name: str) -> JSONResponse:
from litellm.proxy.proxy_server import llm_router
model_ids: Final = [] if llm_router is None else llm_router.get_model_ids(model_name=model_name)
return JSONResponse(content={"model_ids": model_ids}, headers={"x-integration-worker-pid": str(os.getpid())})
async def worker_general_setting(name: str) -> JSONResponse:
from litellm.proxy.proxy_server import proxy_config
return JSONResponse(
content={"value": proxy_config.settings.get(name)}, headers={"x-integration-worker-pid": str(os.getpid())}
)
def configure_readiness() -> None:
from litellm.proxy.proxy_server import app
app.add_api_route("/health/integration-readiness", worker_readiness, include_in_schema=False)
app.add_api_route("/health/integration-deployment-ids", worker_deployment_ids, include_in_schema=False)
app.add_api_route("/health/integration-general-setting", worker_general_setting, include_in_schema=False)