diff --git a/tests/integration/_support/database.py b/tests/integration/_support/database.py index 461cdbda1ee..3ad4f205e12 100644 --- a/tests/integration/_support/database.py +++ b/tests/integration/_support/database.py @@ -1,7 +1,12 @@ import os +import uuid +from collections.abc import Generator +from contextlib import contextmanager from typing import Final, LiteralString +from urllib.parse import urlsplit, urlunsplit import psycopg +from psycopg import sql from psycopg.rows import dict_row from pydantic import JsonValue, TypeAdapter @@ -19,3 +24,14 @@ def read_rows( def write_rows(query: LiteralString, parameters: tuple[str, ...], *, database_url: str | None = None) -> None: with psycopg.connect(database_url or os.environ["DATABASE_URL"]) as connection: connection.execute(query, parameters) + + +@contextmanager +def scratch_database() -> Generator[str]: + name: Final = f"integration_{uuid.uuid4().hex}" + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as admin: + admin.execute(sql.SQL("CREATE DATABASE {}").format(sql.Identifier(name))) + try: + yield urlunsplit(urlsplit(os.environ["DATABASE_URL"])._replace(path=f"/{name}")) + finally: + admin.execute(sql.SQL("DROP DATABASE {} WITH (FORCE)").format(sql.Identifier(name))) diff --git a/tests/integration/observability/test_langfuse_delivery.py b/tests/integration/observability/test_langfuse_delivery.py index 9c0e3e407cb..13be5a95887 100644 --- a/tests/integration/observability/test_langfuse_delivery.py +++ b/tests/integration/observability/test_langfuse_delivery.py @@ -2,13 +2,13 @@ import base64 import json import time import uuid -from collections.abc import Mapping, Sequence +from collections.abc import Sequence from pathlib import Path from typing import Final import yaml from integration._support.client import Gateway, eventually, object_value, string_value -from integration._support.database import read_rows, write_rows +from integration._support.database import read_rows, scratch_database from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, Wire, wire_server from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest @@ -24,6 +24,7 @@ PROMPTS_PATH: Final = "/api/public/v2/prompts/" STOCK_CONFIG: Final = Path("tests/integration/proxy_config.yaml") CONFIG_SECTIONS: Final = ("litellm_settings", "environment_variables") LANGFUSE_ENVIRONMENT: Final = ("LANGFUSE_HOST", "LANGFUSE_PUBLIC_KEY", "LANGFUSE_SECRET_KEY") +INHERITED_ENVIRONMENT: Final = (*LANGFUSE_ENVIRONMENT, "DATABASE_URL_READ_REPLICA") _PROXY_CONFIG: Final = TypeAdapter(dict[str, object]) _SETTINGS: Final = TypeAdapter(dict[str, object]) @@ -88,32 +89,14 @@ def _langfuse_environment(langfuse: Wire) -> dict[str, str]: } -def _config_rows() -> list[dict[str, JsonValue]]: +def _config_rows(database_url: str) -> list[dict[str, JsonValue]]: return read_rows( 'SELECT param_name, param_value FROM "LiteLLM_Config" WHERE param_name IN (%s, %s) ORDER BY param_name', CONFIG_SECTIONS, + database_url=database_url, ) -def _restore_config_rows(snapshot: Sequence[Mapping[str, JsonValue]]) -> None: - saved: Final = {string_value(row["param_name"]): row["param_value"] for row in snapshot} - for section in CONFIG_SECTIONS: - if section not in saved: - write_rows('DELETE FROM "LiteLLM_Config" WHERE param_name = %s', (section,)) - elif saved[section] is None: - write_rows( - 'INSERT INTO "LiteLLM_Config" (param_name, param_value) VALUES (%s, NULL) ' - "ON CONFLICT (param_name) DO UPDATE SET param_value = NULL", - (section,), - ) - else: - write_rows( - 'INSERT INTO "LiteLLM_Config" (param_name, param_value) VALUES (%s, %s::jsonb) ' - "ON CONFLICT (param_name) DO UPDATE SET param_value = EXCLUDED.param_value", - (section, json.dumps(saved[section])), - ) - - def _attribute(entries: Sequence[KeyValue], key: str) -> str | list[str] | None: for entry in entries: if entry.key != key: @@ -227,7 +210,12 @@ def test_langfuse_callback_stored_in_the_db_through_config_update_delivers_the_g provider_secret: Final = "synthetic-provider-secret-" + marker public_key: Final = "pk-lf-db-" + marker secret_key: Final = "sk-lf-db-" + marker - assert "langfuse" not in STOCK_CONFIG.read_text() + stock_settings: Final = _SETTINGS.validate_python( + _PROXY_CONFIG.validate_python(yaml.safe_load(STOCK_CONFIG.read_text()))["litellm_settings"] + ) + assert "langfuse" not in json.dumps( + [stock_settings.get(key) for key in ("callbacks", "success_callback", "failure_callback")] + ) def upstream(request: Request) -> Reply: assert request.headers["authorization"] == f"Bearer {provider_secret}" @@ -238,73 +226,69 @@ def test_langfuse_callback_stored_in_the_db_through_config_update_delivers_the_g return _projects() return Reply(body=b"", content_type="application/x-protobuf") - snapshot: Final = _config_rows() - try: - with ( - wire_server(upstream) as provider, - wire_server(langfuse) as destination, - owned_proxy( - gateway, - tmp_path, - {"LANGFUSE_FLUSH_INTERVAL": "1"}, - remove_environment=LANGFUSE_ENVIRONMENT, - ) as candidate, - candidate.scenario() as scenario, - ): - candidate.post( - "/config/update", - { - "litellm_settings": {"success_callback": ["langfuse"]}, - "environment_variables": { - "LANGFUSE_HOST": destination.url, - "LANGFUSE_PUBLIC_KEY": public_key, - "LANGFUSE_SECRET_KEY": secret_key, - }, + with ( + scratch_database() as scratch_url, + wire_server(upstream) as provider, + wire_server(langfuse) as destination, + owned_proxy( + gateway, + tmp_path, + {"DATABASE_URL": scratch_url, "LANGFUSE_FLUSH_INTERVAL": "1"}, + remove_environment=INHERITED_ENVIRONMENT, + ) as candidate, + candidate.scenario() as scenario, + ): + candidate.post( + "/config/update", + { + "litellm_settings": {"success_callback": ["langfuse"]}, + "environment_variables": { + "LANGFUSE_HOST": destination.url, + "LANGFUSE_PUBLIC_KEY": public_key, + "LANGFUSE_SECRET_KEY": secret_key, }, - ) - model: Final = scenario.model(api_base=provider.url + "/v1", api_key=provider_secret) - body: Final = candidate.post( - "/v1/chat/completions", - { - "model": model, - "messages": [{"role": "user", "content": marker + "-question"}], - "metadata": {"generation_name": marker}, - "cache": {"no-cache": True}, - }, - ) - received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier ones + }, + ) + model: Final = scenario.model(api_base=provider.url + "/v1", api_key=provider_secret) + body: Final = candidate.post( + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": marker + "-question"}], + "metadata": {"generation_name": marker}, + "cache": {"no-cache": True}, + }, + ) + received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier ones - def exported() -> tuple[Span, ...]: - received.extend(destination.drain()) - return tuple(span for span in _spans(received) if span.name == marker) + def exported() -> tuple[Span, ...]: + received.extend(destination.drain()) + return tuple(span for span in _spans(received) if span.name == marker) - spans: Final = eventually(exported, lambda values: len(values) == 1, seconds=20) - posts: Final = tuple(request for request in received if request.method == "POST") - assert {request.target for request in posts} == {TRACES_PATH}, [request.target for request in received] - basic: Final = "Basic " + base64.b64encode(f"{public_key}:{secret_key}".encode()).decode() - for request in posts: - assert request.headers["authorization"] == basic - assert request.headers["content-type"] == "application/x-protobuf" - assert request.headers["x-langfuse-ingestion-version"] == "4" - assert provider_secret.encode() not in request.body - assert candidate.key.encode() not in request.body + spans: Final = eventually(exported, lambda values: len(values) == 1, seconds=20) + posts: Final = tuple(request for request in received if request.method == "POST") + assert {request.target for request in posts} == {TRACES_PATH}, [request.target for request in received] + basic: Final = "Basic " + base64.b64encode(f"{public_key}:{secret_key}".encode()).decode() + for request in posts: + assert request.headers["authorization"] == basic + assert request.headers["content-type"] == "application/x-protobuf" + assert request.headers["x-langfuse-ingestion-version"] == "4" + assert provider_secret.encode() not in request.body + assert candidate.key.encode() not in request.body - attributes: Final = spans[0].attributes - assert _attribute(attributes, "langfuse.observation.type") == "generation" - assert _attribute(attributes, "langfuse.observation.metadata.response_id") == string_value(body["id"]) - assert marker + "-question" in str(_attribute(attributes, "langfuse.observation.input")) - assert marker + "-answer" in str(_attribute(attributes, "langfuse.observation.output")) + attributes: Final = spans[0].attributes + assert _attribute(attributes, "langfuse.observation.type") == "generation" + assert _attribute(attributes, "langfuse.observation.metadata.response_id") == string_value(body["id"]) + assert marker + "-question" in str(_attribute(attributes, "langfuse.observation.input")) + assert marker + "-answer" in str(_attribute(attributes, "langfuse.observation.output")) - stored: Final = {string_value(row["param_name"]): row["param_value"] for row in _config_rows()} - callbacks: Final = TypeAdapter(list[str]).validate_python( - object_value(stored["litellm_settings"]).get("success_callback") or [] - ) - assert "langfuse" in callbacks, stored - assert set(object_value(stored["environment_variables"])) >= set(LANGFUSE_ENVIRONMENT), stored - assert secret_key not in json.dumps(stored["environment_variables"]), stored - finally: - _restore_config_rows(snapshot) - assert _config_rows() == snapshot + stored: Final = {string_value(row["param_name"]): row["param_value"] for row in _config_rows(scratch_url)} + callbacks: Final = TypeAdapter(list[str]).validate_python( + object_value(stored["litellm_settings"]).get("success_callback") or [] + ) + assert "langfuse" in callbacks, stored + assert set(object_value(stored["environment_variables"])) >= set(LANGFUSE_ENVIRONMENT), stored + assert secret_key not in json.dumps(stored["environment_variables"]), stored def test_prompt_fetch_encodes_the_name_retries_a_5xx_once_and_keeps_langfuse_headers_off_the_client( diff --git a/tests/integration/spend/test_cache_and_quota.py b/tests/integration/spend/test_cache_and_quota.py index 1c2b5551855..1cc03838f8d 100644 --- a/tests/integration/spend/test_cache_and_quota.py +++ b/tests/integration/spend/test_cache_and_quota.py @@ -1,27 +1,22 @@ import json -import os import threading import uuid -from collections.abc import Generator from concurrent.futures import ThreadPoolExecutor -from contextlib import ExitStack, contextmanager +from contextlib import ExitStack from hashlib import sha256 from pathlib import Path from typing import Final -from urllib.parse import urlsplit, urlunsplit import httpx -import psycopg import pytest from hypothesis import strategies as st from hypothesis.stateful import RuleBasedStateMachine, rule, run_state_machine_as_test from integration._support.client import Gateway, eventually, string_value -from integration._support.database import read_rows +from integration._support.database import read_rows, scratch_database from integration._support.database_relay import database_relay from integration._support.generation import LIFECYCLE_SETTINGS, bounded_http_requests from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, wire_server -from psycopg import sql @pytest.mark.covers("quota_management.response_cache.generated_sequences_preserve_content_and_accounting") @@ -225,17 +220,6 @@ def test_key_budget_at_boundary_blocks_provider_then_explicit_reset_restores(gat RESET_SWEEP_QUERY: Final = b'"LiteLLM_VerificationToken"."budget_reset_at" < $' -@contextmanager -def scratch_database() -> Generator[str]: - name: Final = f"integration_{uuid.uuid4().hex}" - with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as admin: - admin.execute(sql.SQL("CREATE DATABASE {}").format(sql.Identifier(name))) - try: - yield urlunsplit(urlsplit(os.environ["DATABASE_URL"])._replace(path=f"/{name}")) - finally: - admin.execute(sql.SQL("DROP DATABASE {} WITH (FORCE)").format(sql.Identifier(name))) - - @pytest.mark.covers("quota_management.budget.key.scheduled_reset_survives_transient_db_outage") @pytest.mark.timeout(300) def test_scheduled_budget_reset_reconnects_after_db_transport_failure_and_unblocks_key(