scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

test(integration): let an owned proxy opt into a private database and Redis

Browse source 
Ports the harness part of #47213 (process.py, workers.py, the remaining client.py hunks and its unit tests) so owned_proxy and owned_proxy_process accept state=OwnedState.PRIVATE and boot on a scratch database and their own redis-server, and Scenario cleanup checks read the database the proxy writes to.

Adapted for this line. Main makes private the default and marks 189 test files with shares_worker_state to keep them on the worker's state. This line keeps the worker's state as the default, so no existing test changes behavior and only callers that pass state=OwnedState.PRIVATE go private. The isolated_gateway rewrite and the conftest marker plumbing are left out, and the unit tests cover the shared default plus the explicit private opt-in

(cherry picked from commit fc27dfdc7d)
This commit is contained in:
yuneng-jiang 2026-11-10 15:16:26 -08:00 • committed by Yuneng Jiang
parent 6b79c8166a
commit f0f514b8bd
No known key found for this signature in database
4 changed files with 283 additions and 56 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

24
tests/integration/_support/client.py
View file

@ -35,9 +35,9 @@ def list_value(value: JsonValue) -> Sequence[JsonValue]:
def delete_key_if_present(candidate: Gateway, key: str) -> None:
digest: Final = sha256(key.encode()).hexdigest()
if read_rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)):
if candidate.rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)):
candidate.post("/key/delete", {"keys": [key]})
assert read_rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)) == []
assert candidate.rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)) == []
def eventually(
@ -179,7 +179,7 @@ class Scenario:
def delete_team(self, identity: str) -> None:
self.gateway.post("/team/delete", {"team_ids": [identity]})
assert read_rows('SELECT team_id FROM "LiteLLM_TeamTable" WHERE team_id = %s', (identity,)) == []
assert self.gateway.rows('SELECT team_id FROM "LiteLLM_TeamTable" WHERE team_id = %s', (identity,)) == []
def project(self, team_id: str, **fields: JsonValue) -> str:
created: Final = self.gateway.post(
@ -192,7 +192,9 @@ class Scenario:
def delete_project(self, identity: str) -> None:
response: Final = self.gateway.request("DELETE", "/project/delete", {"project_ids": [identity]})
assert response.status_code == 200, response.text
assert read_rows('SELECT project_id FROM "LiteLLM_ProjectTable" WHERE project_id = %s', (identity,)) == []
assert (
self.gateway.rows('SELECT project_id FROM "LiteLLM_ProjectTable" WHERE project_id = %s', (identity,)) == []
)
def budget(self, **fields: JsonValue) -> str:
created: Final = self.gateway.post("/budget/new", fields)
@ -202,7 +204,7 @@ class Scenario:
def delete_budget(self, identity: str) -> None:
self.gateway.post("/budget/delete", {"id": identity})
assert read_rows('SELECT budget_id FROM "LiteLLM_BudgetTable" WHERE budget_id = %s', (identity,)) == []
assert self.gateway.rows('SELECT budget_id FROM "LiteLLM_BudgetTable" WHERE budget_id = %s', (identity,)) == []
def organization(self, **fields: JsonValue) -> str:
created: Final = self.gateway.post(
@ -216,7 +218,9 @@ class Scenario:
response: Final = self.gateway.request("DELETE", "/organization/delete", {"organization_ids": [identity]})
assert response.status_code == 200, response.text
assert (
read_rows('SELECT organization_id FROM "LiteLLM_OrganizationTable" WHERE organization_id = %s', (identity,))
self.gateway.rows(
'SELECT organization_id FROM "LiteLLM_OrganizationTable" WHERE organization_id = %s', (identity,)
)
== []
)
self.delete_budget(budget_id)
@ -240,7 +244,7 @@ class Scenario:
def delete_user(self, identity: str) -> None:
response: Final = self.gateway.request("POST", "/user/delete", {"user_ids": [identity]})
assert response.status_code == 200 and response.json() == 1, response.text
assert read_rows('SELECT user_id FROM "LiteLLM_UserTable" WHERE user_id = %s', (identity,)) == []
assert self.gateway.rows('SELECT user_id FROM "LiteLLM_UserTable" WHERE user_id = %s', (identity,)) == []
def member(self, team_id: str, role: str = "user") -> str:
"""Create an internal user and add them to ``team_id``; deleting the user later removes the membership."""
@ -251,7 +255,7 @@ class Scenario:
def delete_key(self, token: str) -> None:
self.gateway.post("/key/delete", {"keys": [token]})
hashed: Final = sha256(token.encode()).hexdigest()
assert read_rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE token = %s', (hashed,)) == []
assert self.gateway.rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE token = %s', (hashed,)) == []
info: Final = object_value(self.gateway.get("/key/info", {"key": hashed})["info"])
assert info["status"] == "deleted", f"Deleted key still served as live: {info['status']}"
@ -269,7 +273,9 @@ class Scenario:
lambda served: identity not in served,
seconds=2 * PROXY_CONFIG_RELOAD_INTERVAL_SECONDS,
)
assert read_rows('SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (identity,)) == []
assert (
self.gateway.rows('SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (identity,)) == []
)
def model(
self, *, name: str | None = None, model_info: Mapping[str, JsonValue] | None = None, **parameters: JsonValue

193
tests/integration/_support/process.py
View file

@ -7,9 +7,9 @@ import subprocess
import sys
import time
import uuid
from collections.abc import Generator, Iterable, Iterator, Mapping
from collections.abc import Callable, Generator, Iterable, Iterator, Mapping
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager, suppress
from contextlib import AbstractContextManager, contextmanager, suppress
from dataclasses import dataclass
from enum import Enum
from datetime import UTC, datetime, timedelta
@ -17,7 +17,7 @@ from ipaddress import IPv4Address
from itertools import accumulate
from pathlib import Path
from types import MappingProxyType
from typing import TYPE_CHECKING, Final
from typing import TYPE_CHECKING, Final, TypeAlias
from urllib.parse import urlsplit
import httpx
@ -33,6 +33,8 @@ from tests.integration._support.budget import import_seconds, readiness_for, rea
from tests.integration._support.client import GATEWAY_LIMITS, Gateway, gateway_from_environment
from tests.integration._support.database import scratch_database
from tests.integration._support.ports import allocate_port
from tests.integration._support.redis_process import owned_redis
from tests.integration._support.workers import gateway_profile, replica_urls
if TYPE_CHECKING:
import pytest
@ -43,6 +45,75 @@ LEGACY_MIGRATE_DEPLOY: Final = ("--use_legacy_migration_resolver",)
CONFIGURATION: Final = TypeAdapter(dict[str, object])
_DEPLOYMENT_IDS: Final = TypeAdapter(dict[str, list[str]])
MCP_COVERAGE_CONFIG: Final = Path(__file__).resolve().parents[1] / "mcp_coverage.toml"
OWNED_STATE_VARIABLE: Final = "INTEGRATION_OWNED_STATE"
class OwnedState(Enum):
PRIVATE = "private"
SHARED_WITH_WORKER = "shared_with_worker"
@dataclass(frozen=True, slots=True)
class PrivateState:
database_url: str
environment: Mapping[str, str]
PrivateStateFactory: TypeAlias = Callable[[Path, Mapping[str, str]], AbstractContextManager[PrivateState]]
def boots_private_state(
state: OwnedState | None,
overrides: Mapping[str, str],
remove_environment: tuple[str, ...],
inherited: Mapping[str, str],
) -> bool:
chosen: Final = state or OwnedState(inherited.get(OWNED_STATE_VARIABLE, OwnedState.SHARED_WITH_WORKER.value))
return chosen is OwnedState.PRIVATE and "DATABASE_URL" not in overrides and "DATABASE_URL" not in remove_environment
@contextmanager
def private_state(directory: Path, inherited: Mapping[str, str]) -> Iterator[PrivateState]:
with scratch_database("owned") as database, owned_redis(directory, inherited=inherited) as cache:
replica: Final = replica_urls(database, inherited) if gateway_profile(inherited).replica else None
yield PrivateState(
database,
MappingProxyType(
{
"DATABASE_URL": database if replica is None else replica["INTEGRATION_PROXY_DATABASE_URL"],
**(
{}
if replica is None
else {"DATABASE_URL_READ_REPLICA": replica["INTEGRATION_PROXY_READ_REPLICA_URL"]}
),
"REDIS_HOST": cache.host,
"REDIS_PORT": str(cache.port),
}
),
)
@dataclass(frozen=True, slots=True)
class OwnedEnvironment:
overrides: Mapping[str, str]
database_url: str | None
@contextmanager
def owned_environment(
directory: Path,
overrides: Mapping[str, str],
remove_environment: tuple[str, ...],
state: OwnedState | None,
inherited: Mapping[str, str],
private: PrivateStateFactory = private_state,
) -> Iterator[OwnedEnvironment]:
if not boots_private_state(state, overrides, remove_environment, inherited):
yield OwnedEnvironment(overrides, None)
return
with private(directory, inherited) as booted:
kept: Final = {name: value for name, value in booted.environment.items() if name not in remove_environment}
yield OwnedEnvironment(MappingProxyType({**kept, **overrides}), booted.database_url)
def proxy_database_environment(inherited: Mapping[str, str] = os.environ) -> Mapping[str, str]:
@ -368,6 +439,7 @@ def owned_proxy(
database_setup: tuple[str, ...] = DB_PUSH,
coverage_file: Path | None = None,
inherited: Mapping[str, str] = os.environ,
state: OwnedState | None = None,
) -> Iterator[Gateway]:
with owned_proxy_process(
gateway,
@ -379,6 +451,7 @@ def owned_proxy(
database_setup=database_setup,
coverage_file=coverage_file,
inherited=inherited,
state=state,
) as owned:
yield owned.gateway
@ -634,6 +707,32 @@ def _proxy_root() -> Path:
return Path(os.environ.get("INTEGRATION_PROXY_ROOT") or Path(__file__).resolve().parents[3])
def _configuration_copy(config: Path, directory: Path, prefix: str) -> Path:
checkouts: Final = (_harness_root(), _proxy_root().resolve())
inside_checkout: Final = any(config.resolve().is_relative_to(checkout) for checkout in checkouts)
target: Final = directory if inside_checkout else config.parent
target.mkdir(parents=True, exist_ok=True)
return target / f"{prefix}-{uuid.uuid4().hex}.yaml"
def private_configuration(config: Path, directory: Path) -> Path:
loaded: Final[object] = yaml.safe_load(config.read_text())
if not isinstance(loaded, dict):
return config
settings: Final = CONFIGURATION.validate_python(loaded)
router: Final = settings.get("router_settings") or {}
if not isinstance(router, dict) or "num_retries" in router:
return config
destination: Final = _configuration_copy(config, directory, "private-config")
destination.write_text(
yaml.safe_dump(
{**settings, "router_settings": {**CONFIGURATION.validate_python(router), "num_retries": 0}},
sort_keys=False,
)
)
return destination
def worker_configuration(config: Path, directory: Path, inherited: Mapping[str, str] = os.environ) -> Path:
if inherited.get("INTEGRATION_WORKER_ISOLATION") != "1":
return config
@ -644,11 +743,7 @@ def worker_configuration(config: Path, directory: Path, inherited: Mapping[str,
return f"http://127.0.0.1:{inherited['INTEGRATION_PROVIDER_PORT']}"
contents: Final = re.sub(r"http://127[.]0[.]0[.]1:(8190|8191)(?![0-9])", endpoint, config.read_text())
checkouts: Final = (_harness_root(), _proxy_root().resolve())
inside_checkout: Final = any(config.resolve().is_relative_to(checkout) for checkout in checkouts)
target: Final = directory if inside_checkout else config.parent
target.mkdir(parents=True, exist_ok=True)
destination: Final = target / f"worker-config-{uuid.uuid4().hex}.yaml"
destination: Final = _configuration_copy(config, directory, "worker-config")
loaded: Final[object] = yaml.safe_load(contents)
if not isinstance(loaded, dict):
destination.write_text(contents)
@ -753,42 +848,58 @@ def owned_proxy_process(
healthcheck: WorkerHealthcheck = WorkerHealthcheck.STRICT,
inherited: Mapping[str, str] = os.environ,
client_timeout: float = 15,
state: OwnedState | None = None,
) -> Iterator[OwnedProxy]:
assert "--timeout_worker_healthcheck" not in extra_arguments, "Choose the worker healthcheck with `healthcheck`"
root: Final = _proxy_root()
environment: Final = _proxy_environment(
gateway,
{**overrides, **({"COVERAGE_FILE": str(coverage_file)} if coverage_file else {})},
remove_environment,
inherited,
)
output: Final = Path(inherited.get("INTEGRATION_RESULTS_DIR", str(directory)))
output.mkdir(parents=True, exist_ok=True)
command: Final = (
sys.executable,
"-m",
*(("coverage", "run", f"--rcfile={MCP_COVERAGE_CONFIG}", "-m") if coverage_file else ()),
"integration._support.proxy",
"--config",
str(worker_configuration(config or root / "tests/integration/proxy_config.yaml", directory, inherited)),
"--host",
"127.0.0.1",
"--num_workers",
str(workers),
"--timeout_worker_healthcheck",
str(healthcheck_seconds(healthcheck)),
*database_setup,
*extra_arguments,
)
launch: Final = _launch_until_ready(command, root, environment, output, workers)
process: Final = launch.process
try:
with httpx.Client(
base_url=f"http://127.0.0.1:{launch.port}", timeout=client_timeout, trust_env=False, limits=GATEWAY_LIMITS
) as client:
yield OwnedProxy(Gateway(client, gateway.key, gateway.upstream_url), process, launch.log, healthcheck)
finally:
_stop_launch(launch)
with owned_environment(directory, overrides, remove_environment, state, inherited) as owned:
configured: Final = worker_configuration(
config or root / "tests/integration/proxy_config.yaml", directory, inherited
)
owned_configuration: Final = (
configured if owned.database_url is None else private_configuration(configured, directory)
)
environment: Final = _proxy_environment(
gateway,
{**owned.overrides, **({"COVERAGE_FILE": str(coverage_file)} if coverage_file else {})},
remove_environment,
inherited,
)
output: Final = Path(inherited.get("INTEGRATION_RESULTS_DIR", str(directory)))
output.mkdir(parents=True, exist_ok=True)
command: Final = (
sys.executable,
"-m",
*(("coverage", "run", f"--rcfile={MCP_COVERAGE_CONFIG}", "-m") if coverage_file else ()),
"integration._support.proxy",
"--config",
str(owned_configuration),
"--host",
"127.0.0.1",
"--num_workers",
str(workers),
"--timeout_worker_healthcheck",
str(healthcheck_seconds(healthcheck)),
*database_setup,
*extra_arguments,
)
launch: Final = _launch_until_ready(command, root, environment, output, workers)
process: Final = launch.process
try:
with httpx.Client(
base_url=f"http://127.0.0.1:{launch.port}",
timeout=client_timeout,
trust_env=False,
limits=GATEWAY_LIMITS,
) as client:
yield OwnedProxy(
Gateway(client, gateway.key, gateway.upstream_url, owned.database_url),
process,
launch.log,
healthcheck,
)
finally:
_stop_launch(launch)
@contextmanager

8
tests/integration/_support/workers.py
View file

@ -86,9 +86,9 @@ def worker_endpoints() -> Mapping[str, str]:
)
def _replica_urls(database: str) -> Mapping[str, str]:
writer: Final = urlsplit(os.environ["INTEGRATION_PROXY_DATABASE_URL"])
reader: Final = urlsplit(os.environ["INTEGRATION_PROXY_READ_REPLICA_URL"])
def replica_urls(database: str, inherited: Mapping[str, str] = os.environ) -> Mapping[str, str]:
writer: Final = urlsplit(inherited["INTEGRATION_PROXY_DATABASE_URL"])
reader: Final = urlsplit(inherited["INTEGRATION_PROXY_READ_REPLICA_URL"])
database_path: Final = urlsplit(database).path
assert writer.username is not None and reader.username is not None
with psycopg.connect(database, autocommit=True) as admin:
@ -164,7 +164,7 @@ def worker_services(
"INTEGRATION_PROXY_READ_REPLICA_URL": "",
"REDIS_HOST": cache.host,
"REDIS_PORT": str(cache.port),
**(_replica_urls(database) if replica else {}),
**(replica_urls(database, inherited) if replica else {}),
}
)

114
tests/unit/integration_support/test_process.py
View file

@ -8,11 +8,11 @@ import subprocess
import sys
import threading
import time
from collections.abc import Iterator
from collections.abc import Iterator, Mapping
from contextlib import contextmanager
from functools import reduce
from pathlib import Path
from types import ModuleType
from types import MappingProxyType, ModuleType
from typing import Final
from unittest.mock import MagicMock
@ -25,13 +25,20 @@ from pydantic import TypeAdapter
from litellm.proxy.types_utils.utils import get_instance_fn
from tests.integration._support.client import Gateway, eventually
from tests.integration._support.process import (
OWNED_STATE_VARIABLE,
OwnedProxy,
OwnedState,
PrivateState,
WorkerHealthcheck,
WorkerKill,
_is_gone,
_proxy_environment,
boots_private_state,
healthcheck_seconds,
kill_stray_process,
owned_environment,
owned_proxy_process,
private_configuration,
replacement_pid,
stop_process_group,
worker_configuration,
@ -385,6 +392,109 @@ def test_owned_runtime_configuration_preserves_custom_worker_startup_hooks(
)
WORKER_STATE: Final = MappingProxyType(
{"DATABASE_URL": "postgresql://admin@127.0.0.1/worker", "REDIS_HOST": "127.0.0.1", "REDIS_PORT": "6379"}
)
PRIVATE_DATABASE: Final = "postgresql://admin@127.0.0.1/owned_private"
@contextmanager
def _private_state(directory: Path, inherited: Mapping[str, str]) -> Iterator[PrivateState]:
yield PrivateState(
PRIVATE_DATABASE,
MappingProxyType({"DATABASE_URL": PRIVATE_DATABASE, "REDIS_HOST": "127.0.0.2", "REDIS_PORT": "6390"}),
)
@pytest.mark.parametrize(
("state", "overrides", "remove_environment", "marked", "private"),
(
(None, {}, (), False, False),
(OwnedState.PRIVATE, {}, (), False, True),
(OwnedState.PRIVATE, {"DATABASE_URL": "postgresql://admin@127.0.0.1/scratch"}, (), False, False),
(OwnedState.PRIVATE, {}, ("DATABASE_URL",), False, False),
(OwnedState.SHARED_WITH_WORKER, {}, (), False, False),
(None, {}, (), True, False),
(OwnedState.PRIVATE, {}, (), True, True),
),
ids=(
"default",
"private",
"explicit-database",
"stripped-database",
"shared",
"marked-module",
"private-in-marked-module",
),
)
def test_an_owned_proxy_boots_private_state_only_when_asked_and_without_its_own_database(
state: OwnedState | None,
overrides: Mapping[str, str],
remove_environment: tuple[str, ...],
marked: bool,
private: bool,
) -> None:
marker: Final = {OWNED_STATE_VARIABLE: OwnedState.SHARED_WITH_WORKER.value} if marked else {}
inherited: Final = {**WORKER_STATE, **marker}
assert boots_private_state(state, overrides, remove_environment, inherited) is private
def test_a_private_owned_proxy_runs_on_its_own_database_and_redis(tmp_path: Path) -> None:
with (
httpx.Client(base_url="http://127.0.0.1:1", trust_env=False) as client,
owned_environment(tmp_path, {}, (), OwnedState.PRIVATE, WORKER_STATE, _private_state) as owned,
):
environment: Final = _proxy_environment(
Gateway(client, "owned-test-key", "http://127.0.0.1:1"), owned.overrides, (), WORKER_STATE
)
assert environment["DATABASE_URL"] == PRIVATE_DATABASE != WORKER_STATE["DATABASE_URL"]
assert (environment["REDIS_HOST"], environment["REDIS_PORT"]) == ("127.0.0.2", "6390")
assert owned.database_url == PRIVATE_DATABASE
@pytest.mark.parametrize(
("overrides", "remove_environment", "redis"),
(
({}, ("REDIS_HOST", "REDIS_PORT"), (None, None)),
({"REDIS_HOST": "redis.caller", "REDIS_PORT": "7000"}, (), ("redis.caller", "7000")),
),
ids=("stripped", "caller-owned"),
)
def test_a_private_owned_proxy_leaves_redis_to_a_caller_that_chose_it(
tmp_path: Path,
overrides: Mapping[str, str],
remove_environment: tuple[str, ...],
redis: tuple[str | None, str | None],
) -> None:
with (
httpx.Client(base_url="http://127.0.0.1:1", trust_env=False) as client,
owned_environment(
tmp_path, overrides, remove_environment, OwnedState.PRIVATE, WORKER_STATE, _private_state
) as owned,
):
environment: Final = _proxy_environment(
Gateway(client, "owned-test-key", "http://127.0.0.1:1"), owned.overrides, remove_environment, WORKER_STATE
)
assert environment["DATABASE_URL"] == PRIVATE_DATABASE
assert (environment.get("REDIS_HOST"), environment.get("REDIS_PORT")) == redis
@pytest.mark.parametrize(
("router_settings", "retries"),
(({"allowed_fails": 3}, 0), ({"num_retries": 2}, 2), (None, 0)),
ids=("inherits-worker-default", "keeps-explicit", "no-router-settings"),
)
def test_a_private_owned_proxy_keeps_the_worker_retry_default_unless_its_config_sets_one(
tmp_path: Path, router_settings: dict[str, int] | None, retries: int
) -> None:
source: Final = tmp_path / "source.yaml"
router: Final = {"router_settings": router_settings} if router_settings else {}
source.write_text(yaml.safe_dump({"model_list": [], **router}))
configured: Final = yaml.safe_load(private_configuration(source, tmp_path / "proxy").read_text())
assert configured["router_settings"]["num_retries"] == retries
assert configured["model_list"] == []
_SUPERVISOR: Final = """import multiprocessing
import time