diff --git a/tests/integration/_support/client.py b/tests/integration/_support/client.py index fe1db22987d..23d63d100b9 100644 --- a/tests/integration/_support/client.py +++ b/tests/integration/_support/client.py @@ -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 diff --git a/tests/integration/_support/process.py b/tests/integration/_support/process.py index d5ce6f916a2..70b701e52b0 100644 --- a/tests/integration/_support/process.py +++ b/tests/integration/_support/process.py @@ -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 diff --git a/tests/integration/_support/workers.py b/tests/integration/_support/workers.py index e75b0b3c6bb..f1958ef6990 100644 --- a/tests/integration/_support/workers.py +++ b/tests/integration/_support/workers.py @@ -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 {}), } ) diff --git a/tests/unit/integration_support/test_process.py b/tests/unit/integration_support/test_process.py index 7741de93a80..585974d3dd5 100644 --- a/tests/unit/integration_support/test_process.py +++ b/tests/unit/integration_support/test_process.py @@ -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