scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-27 01:22:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 82

test(integration): run the Langfuse DB-callback test on its own scratch database (#43288)

Browse source 
* test(integration): run the Langfuse DB-callback test on its own scratch database

The test from #43282 wrote success_callback=langfuse and the LANGFUSE_* env
into the shared integration LiteLLM_Config. The suite's long-running gateway
reloads that table and only ever adds callbacks, so it kept exporting to the
test's closed Langfuse fake for the rest of the shard even after the rows were
restored. The owned proxy now gets a scratch database, which also removes the
snapshot/restore code. scratch_database moves into _support/database.py so
test_cache_and_quota and this test share one copy, and the stock-config guard
now checks the callback settings instead of the raw YAML text.

* test(integration): include failure_callback in the stock-config Langfuse guard
This commit is contained in:
yuneng-jiang 2026-09-26 00:15:56 -07:00 • committed by GitHub
parent a942c343ab
commit dd63637322
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 87 additions and 103 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

16
tests/integration/_support/database.py
View file

@ -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)))

154
tests/integration/observability/test_langfuse_delivery.py
View file

@ -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(

20
tests/integration/spend/test_cache_and_quota.py
View file

@ -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(