mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-11 03:38:38 +00:00
test(spend): rework audit cells for review findings
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
kerry
parent
f9282fd28d
commit
3d91b1dcf1
2 changed files with 203 additions and 212 deletions
|
|
@ -6,7 +6,9 @@ from hashlib import sha256
|
|||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import anthropic
|
||||
import httpx
|
||||
import openai
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import Gateway, eventually
|
||||
|
|
@ -29,46 +31,18 @@ SENT_HEADERS: Final = {"user-agent": "claude-cli/2.0.0", "x-tenant-id": "tenant-
|
|||
EXPECTED_TAGS: Final = ["User-Agent: claude-cli", "User-Agent: claude-cli/2.0.0", "x-tenant-id: tenant-a"]
|
||||
|
||||
|
||||
def _respond(request: Request) -> Reply:
|
||||
assert request.method == "POST" and request.target == "/v1/messages", request.target
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": f"msg_{uuid.uuid4().hex}",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": MODEL,
|
||||
"content": [{"type": "text", "text": "tagged"}],
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 10, "output_tokens": 2},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
|
||||
def _request_tags(key: str) -> list[list[str]]:
|
||||
rows: Final = read_rows(
|
||||
'SELECT request_tags FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (sha256(key.encode()).hexdigest(),)
|
||||
)
|
||||
return [
|
||||
json.loads(row["request_tags"]) if isinstance(row["request_tags"], str) else row["request_tags"] for row in rows
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"route", [pytest.param("/anthropic/v1/messages", id="passthrough"), pytest.param("/v1/messages", id="unified")]
|
||||
)
|
||||
def test_header_derived_spend_tags_are_recorded_on_anthropic_messages_routes(
|
||||
gateway: Gateway, tmp_path: Path, route: str
|
||||
) -> None:
|
||||
with wire_server(_respond) as wire:
|
||||
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
config["litellm_settings"]["extra_spend_tag_headers"] = ["x-tenant-id"]
|
||||
path: Final = tmp_path / "spend-tag-headers.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
environment: Final = {"ANTHROPIC_API_BASE": wire.url, "ANTHROPIC_API_KEY": "synthetic-anthropic-key"}
|
||||
with owned_proxy(gateway, tmp_path, environment, config=path) as candidate, candidate.scenario() as scenario:
|
||||
with wire_server(provider_reply) as wire:
|
||||
config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
):
|
||||
model: Final = scenario.model(
|
||||
model=f"anthropic/{MODEL}", api_base=wire.url, api_key="synthetic-anthropic-key"
|
||||
)
|
||||
|
|
@ -86,30 +60,21 @@ def test_header_derived_spend_tags_are_recorded_on_anthropic_messages_routes(
|
|||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert len(wire.drain()) == 1
|
||||
assert eventually(lambda: _request_tags(key), lambda tags: len(tags) == 1, seconds=70) == [EXPECTED_TAGS]
|
||||
assert eventually(lambda: tags_by_key(key), lambda tags: len(tags) == 1, seconds=70) == [EXPECTED_TAGS]
|
||||
|
||||
|
||||
def _base_url(candidate: Gateway) -> str:
|
||||
return str(candidate.client.base_url).rstrip("/")
|
||||
|
||||
|
||||
def _spend_count() -> int:
|
||||
return read_rows('SELECT count(*) AS n FROM "LiteLLM_SpendLogs"', ())[0]["n"]
|
||||
|
||||
|
||||
def _owned_config(tmp_path: Path, mutations: dict) -> Path:
|
||||
return write_config(tmp_path, mutations)
|
||||
|
||||
|
||||
UA_TAG: Final = "User-Agent: claude-cli/2.0.0"
|
||||
UA_FAMILY_TAG: Final = "User-Agent: claude-cli"
|
||||
TENANT_TAG: Final = "x-tenant-id: tenant-a"
|
||||
|
||||
|
||||
# H2: pass-through streaming anthropic request records the header tags
|
||||
def test_pass_through_anthropic_stream_records_header_tags(gateway: Gateway, tmp_path: Path) -> None:
|
||||
with wire_server(provider_reply) as wire:
|
||||
config: Final = _owned_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
|
|
@ -133,12 +98,9 @@ def test_pass_through_anthropic_stream_records_header_tags(gateway: Gateway, tmp
|
|||
assert eventually(lambda: tags_by_key(key), lambda tags: len(tags) == 1, seconds=70) == [EXPECTED_TAGS]
|
||||
|
||||
|
||||
# H3: Anthropic SDK non-stream call through the pass-through route
|
||||
def test_pass_through_anthropic_sdk_records_header_tags(gateway: Gateway, tmp_path: Path) -> None:
|
||||
import anthropic
|
||||
|
||||
with wire_server(provider_reply) as wire:
|
||||
config: Final = _owned_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
|
|
@ -157,12 +119,9 @@ def test_pass_through_anthropic_sdk_records_header_tags(gateway: Gateway, tmp_pa
|
|||
]
|
||||
|
||||
|
||||
# H4: Anthropic SDK streaming call through the pass-through route
|
||||
def test_pass_through_anthropic_sdk_stream_records_header_tags(gateway: Gateway, tmp_path: Path) -> None:
|
||||
import anthropic
|
||||
|
||||
with wire_server(provider_reply) as wire:
|
||||
config: Final = _owned_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
|
|
@ -182,11 +141,10 @@ def test_pass_through_anthropic_sdk_stream_records_header_tags(gateway: Gateway,
|
|||
]
|
||||
|
||||
|
||||
# H5: openai pass-through route records the header tags
|
||||
@pytest.mark.parametrize("stream", [pytest.param(False, id="sync"), pytest.param(True, id="stream")])
|
||||
def test_pass_through_openai_chat_records_header_tags(gateway: Gateway, tmp_path: Path, stream: bool) -> None:
|
||||
with wire_server(provider_reply) as wire:
|
||||
config: Final = _owned_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
|
|
@ -209,13 +167,10 @@ def test_pass_through_openai_chat_records_header_tags(gateway: Gateway, tmp_path
|
|||
assert eventually(lambda: tags_by_key(key), lambda tags: len(tags) == 1, seconds=70) == [EXPECTED_TAGS]
|
||||
|
||||
|
||||
# H6: OpenAI SDK sync and stream calls through the pass-through route
|
||||
@pytest.mark.parametrize("stream", [pytest.param(False, id="sync"), pytest.param(True, id="stream")])
|
||||
def test_pass_through_openai_sdk_records_header_tags(gateway: Gateway, tmp_path: Path, stream: bool) -> None:
|
||||
import openai
|
||||
|
||||
with wire_server(provider_reply) as wire:
|
||||
config: Final = _owned_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
|
|
@ -240,10 +195,9 @@ def test_pass_through_openai_sdk_records_header_tags(gateway: Gateway, tmp_path:
|
|||
]
|
||||
|
||||
|
||||
# H7: gemini pass-through route records the header tags
|
||||
def test_pass_through_gemini_records_header_tags(gateway: Gateway, tmp_path: Path) -> None:
|
||||
with wire_server(provider_reply) as wire:
|
||||
config: Final = _owned_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
|
|
@ -261,10 +215,9 @@ def test_pass_through_gemini_records_header_tags(gateway: Gateway, tmp_path: Pat
|
|||
assert eventually(lambda: tags_by_key(key), lambda tags: len(tags) == 1, seconds=70) == [EXPECTED_TAGS]
|
||||
|
||||
|
||||
# H8: user-defined pass-through endpoint records the header tags
|
||||
def test_custom_pass_through_endpoint_records_header_tags(gateway: Gateway, tmp_path: Path) -> None:
|
||||
with wire_server(provider_reply) as wire:
|
||||
config: Final = _owned_config(
|
||||
config: Final = write_config(
|
||||
tmp_path,
|
||||
{
|
||||
"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]},
|
||||
|
|
@ -304,12 +257,9 @@ def test_custom_pass_through_endpoint_records_header_tags(gateway: Gateway, tmp_
|
|||
assert eventually(lambda: tags_by_key(key), lambda tags: len(tags) == 1, seconds=70) == [EXPECTED_TAGS]
|
||||
|
||||
|
||||
# H9: async OpenAI SDK streaming call through the pass-through route
|
||||
def test_pass_through_openai_async_sdk_stream_records_header_tags(gateway: Gateway, tmp_path: Path) -> None:
|
||||
import openai
|
||||
|
||||
with wire_server(provider_reply) as wire:
|
||||
config: Final = _owned_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
|
|
@ -338,10 +288,9 @@ def test_pass_through_openai_async_sdk_stream_records_header_tags(gateway: Gatew
|
|||
]
|
||||
|
||||
|
||||
# H10: unified chat completions control records the header tags
|
||||
def test_unified_chat_completions_records_header_tags(gateway: Gateway, tmp_path: Path) -> None:
|
||||
with wire_server(provider_reply) as wire:
|
||||
config: Final = _owned_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
|
|
@ -363,13 +312,10 @@ def test_unified_chat_completions_records_header_tags(gateway: Gateway, tmp_path
|
|||
]
|
||||
|
||||
|
||||
# H11: unified /v1/responses via the OpenAI SDK, sync and stream
|
||||
@pytest.mark.parametrize("stream", [pytest.param(False, id="sync"), pytest.param(True, id="stream")])
|
||||
def test_unified_responses_records_header_tags(gateway: Gateway, tmp_path: Path, stream: bool) -> None:
|
||||
import openai
|
||||
|
||||
with wire_server(provider_reply) as wire:
|
||||
config: Final = _owned_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
|
|
@ -390,10 +336,9 @@ def test_unified_responses_records_header_tags(gateway: Gateway, tmp_path: Path,
|
|||
assert eventually(lambda: tags_by_key(key), lambda tags: len(tags) == 1, seconds=70) == [EXPECTED_TAGS]
|
||||
|
||||
|
||||
# H12: a unified cache-hit twin records the same tags on both spend rows
|
||||
def test_unified_cache_hit_twin_records_header_tags(gateway: Gateway, tmp_path: Path) -> None:
|
||||
with wire_server(provider_reply) as wire:
|
||||
config: Final = _owned_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
|
|
@ -412,13 +357,12 @@ def test_unified_cache_hit_twin_records_header_tags(gateway: Gateway, tmp_path:
|
|||
]
|
||||
|
||||
|
||||
# H13: generic_api callback sink sees the same request_tags as the spend row
|
||||
def test_pass_through_tags_reach_generic_api_sink(gateway: Gateway, tmp_path: Path) -> None:
|
||||
def sink(request: Request) -> Reply:
|
||||
return Reply()
|
||||
|
||||
with wire_server(provider_reply) as wire, wire_server(sink) as endpoint:
|
||||
config: Final = _owned_config(
|
||||
config: Final = write_config(
|
||||
tmp_path,
|
||||
{
|
||||
"litellm_settings": {
|
||||
|
|
@ -470,10 +414,9 @@ def test_pass_through_tags_reach_generic_api_sink(gateway: Gateway, tmp_path: Pa
|
|||
]
|
||||
|
||||
|
||||
# H14: LiteLLM_DailyTagSpend accrues each header-derived tag
|
||||
def test_pass_through_tags_accrue_daily_tag_spend(gateway: Gateway, tmp_path: Path) -> None:
|
||||
with wire_server(provider_reply) as wire:
|
||||
config: Final = _owned_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
|
|
@ -506,7 +449,6 @@ def test_pass_through_tags_accrue_daily_tag_spend(gateway: Gateway, tmp_path: Pa
|
|||
assert {row["tag"] for row in rows} == set(EXPECTED_TAGS)
|
||||
|
||||
|
||||
# S1: extra_spend_tag_headers unset records only the user-agent tags on both routes
|
||||
@pytest.mark.parametrize(
|
||||
"route", [pytest.param("/anthropic/v1/messages", id="passthrough"), pytest.param("/v1/messages", id="unified")]
|
||||
)
|
||||
|
|
@ -514,7 +456,7 @@ def test_header_tags_without_extra_spend_tag_headers_record_user_agent_only(
|
|||
gateway: Gateway, tmp_path: Path, route: str
|
||||
) -> None:
|
||||
with wire_server(provider_reply) as wire:
|
||||
config: Final = _owned_config(tmp_path, {})
|
||||
config: Final = write_config(tmp_path, {})
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
|
|
@ -542,13 +484,12 @@ def test_header_tags_without_extra_spend_tag_headers_record_user_agent_only(
|
|||
]
|
||||
|
||||
|
||||
# S2: pass-through request without the headers records no tags
|
||||
@pytest.mark.parametrize(
|
||||
"route", [pytest.param("/anthropic/v1/messages", id="passthrough"), pytest.param("/v1/messages", id="unified")]
|
||||
)
|
||||
def test_routes_without_headers_record_no_tags(gateway: Gateway, tmp_path: Path, route: str) -> None:
|
||||
with wire_server(provider_reply) as wire:
|
||||
config: Final = _owned_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
|
|
@ -583,13 +524,12 @@ def test_routes_without_headers_record_no_tags(gateway: Gateway, tmp_path: Path,
|
|||
assert eventually(lambda: tags_by_id(request_id), lambda tags: len(tags) == 1, seconds=70) == [[]]
|
||||
|
||||
|
||||
# S3: disable_add_user_agent_to_request_tags keeps only the extra header tags
|
||||
@pytest.mark.parametrize(
|
||||
"route", [pytest.param("/anthropic/v1/messages", id="passthrough"), pytest.param("/v1/messages", id="unified")]
|
||||
)
|
||||
def test_disabled_user_agent_keeps_only_extra_header_tags(gateway: Gateway, tmp_path: Path, route: str) -> None:
|
||||
with wire_server(provider_reply) as wire:
|
||||
config: Final = _owned_config(
|
||||
config: Final = write_config(
|
||||
tmp_path,
|
||||
{
|
||||
"litellm_settings": {
|
||||
|
|
@ -624,13 +564,12 @@ def test_disabled_user_agent_keeps_only_extra_header_tags(gateway: Gateway, tmp_
|
|||
]
|
||||
|
||||
|
||||
# S4: a configured header the client never sends contributes no tag
|
||||
@pytest.mark.parametrize(
|
||||
"route", [pytest.param("/anthropic/v1/messages", id="passthrough"), pytest.param("/v1/messages", id="unified")]
|
||||
)
|
||||
def test_unsent_configured_header_contributes_no_tag(gateway: Gateway, tmp_path: Path, route: str) -> None:
|
||||
with wire_server(provider_reply) as wire:
|
||||
config: Final = _owned_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-never-sent"]}})
|
||||
config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-never-sent"]}})
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
|
|
@ -657,13 +596,12 @@ def test_unsent_configured_header_contributes_no_tag(gateway: Gateway, tmp_path:
|
|||
]
|
||||
|
||||
|
||||
# S5: httpx default user-agent is recorded when the client sends none
|
||||
@pytest.mark.parametrize(
|
||||
"route", [pytest.param("/anthropic/v1/messages", id="passthrough"), pytest.param("/v1/messages", id="unified")]
|
||||
)
|
||||
def test_default_httpx_user_agent_is_recorded(gateway: Gateway, tmp_path: Path, route: str) -> None:
|
||||
with wire_server(provider_reply) as wire:
|
||||
config: Final = _owned_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
|
|
@ -691,16 +629,14 @@ def test_default_httpx_user_agent_is_recorded(gateway: Gateway, tmp_path: Path,
|
|||
]
|
||||
|
||||
|
||||
# S6: unauthenticated requests return 401 and write no spend row
|
||||
@pytest.mark.parametrize("route", [pytest.param("/anthropic/v1/messages", id="passthrough")])
|
||||
def test_unauthenticated_request_writes_no_spend_row(gateway: Gateway, tmp_path: Path, route: str) -> None:
|
||||
def test_unauthenticated_pass_through_writes_untagged_spend_row(gateway: Gateway, tmp_path: Path, route: str) -> None:
|
||||
with wire_server(provider_reply) as wire:
|
||||
config: Final = _owned_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
):
|
||||
before: Final = _spend_count()
|
||||
anonymous_before: Final = len(
|
||||
read_rows("SELECT request_tags FROM \"LiteLLM_SpendLogs\" WHERE api_key IS NULL OR api_key=''", ())
|
||||
)
|
||||
|
|
@ -740,16 +676,14 @@ def test_unauthenticated_request_writes_no_spend_row(gateway: Gateway, tmp_path:
|
|||
assert eventually(lambda: tags_by_id(control.json()["id"]), lambda tags: len(tags) == 1, seconds=70) == [
|
||||
EXPECTED_TAGS
|
||||
]
|
||||
assert _spend_count() == before + 2
|
||||
|
||||
|
||||
# S7: an upstream 400 surfaces the same status and its spend row records the tags
|
||||
@pytest.mark.parametrize(
|
||||
"route", [pytest.param("/anthropic/v1/messages", id="passthrough"), pytest.param("/v1/messages", id="unified")]
|
||||
)
|
||||
def test_upstream_failure_still_records_header_tags(gateway: Gateway, tmp_path: Path, route: str) -> None:
|
||||
with wire_server(provider_reply) as wire:
|
||||
config: Final = _owned_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
|
|
@ -777,7 +711,6 @@ def test_upstream_failure_still_records_header_tags(gateway: Gateway, tmp_path:
|
|||
assert eventually(lambda: tags_by_key(key), lambda tags: len(tags) == 1, seconds=70) == [expected]
|
||||
|
||||
|
||||
# S8: null and empty extra_spend_tag_headers behave like unset
|
||||
@pytest.mark.parametrize("extra", [pytest.param(None, id="null"), pytest.param([], id="empty")])
|
||||
@pytest.mark.parametrize(
|
||||
"route", [pytest.param("/anthropic/v1/messages", id="passthrough"), pytest.param("/v1/messages", id="unified")]
|
||||
|
|
@ -786,7 +719,7 @@ def test_null_and_empty_extra_spend_tag_headers_record_user_agent_only(
|
|||
gateway: Gateway, tmp_path: Path, route: str, extra: object
|
||||
) -> None:
|
||||
with wire_server(provider_reply) as wire:
|
||||
config: Final = _owned_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": extra}})
|
||||
config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": extra}})
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
|
|
@ -813,10 +746,9 @@ def test_null_and_empty_extra_spend_tag_headers_record_user_agent_only(
|
|||
]
|
||||
|
||||
|
||||
# S9: pass-through matches the configured header name case-insensitively, unified is case-sensitive
|
||||
def test_configured_header_case_differs_between_routes(gateway: Gateway, tmp_path: Path) -> None:
|
||||
with wire_server(provider_reply) as wire:
|
||||
config: Final = _owned_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["X-Tenant-Id"]}})
|
||||
config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["X-Tenant-Id"]}})
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
|
|
@ -854,14 +786,13 @@ def test_configured_header_case_differs_between_routes(gateway: Gateway, tmp_pat
|
|||
assert unified_tags == [UA_FAMILY_TAG, UA_TAG], rows
|
||||
|
||||
|
||||
# E1: a 5KB header value is stored verbatim
|
||||
@pytest.mark.parametrize(
|
||||
"route", [pytest.param("/anthropic/v1/messages", id="passthrough"), pytest.param("/v1/messages", id="unified")]
|
||||
)
|
||||
def test_large_header_value_is_stored_verbatim(gateway: Gateway, tmp_path: Path, route: str) -> None:
|
||||
big: Final = "x" * 5000
|
||||
with wire_server(provider_reply) as wire:
|
||||
config: Final = _owned_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
|
|
@ -892,10 +823,9 @@ def test_large_header_value_is_stored_verbatim(gateway: Gateway, tmp_path: Path,
|
|||
]
|
||||
|
||||
|
||||
# E2: duplicate configured headers record first-value on pass-through, last-value on unified
|
||||
def test_duplicate_header_values_follow_carrier_semantics(gateway: Gateway, tmp_path: Path) -> None:
|
||||
with wire_server(provider_reply) as wire:
|
||||
config: Final = _owned_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
|
|
@ -926,13 +856,12 @@ def test_duplicate_header_values_follow_carrier_semantics(gateway: Gateway, tmp_
|
|||
assert tags_by_id(unified.json()["id"])[0] == [UA_FAMILY_TAG, UA_TAG, "x-tenant-id: t2"]
|
||||
|
||||
|
||||
# E4: x-litellm-tags and header-derived tags land together
|
||||
@pytest.mark.parametrize(
|
||||
"route", [pytest.param("/anthropic/v1/messages", id="passthrough"), pytest.param("/v1/messages", id="unified")]
|
||||
)
|
||||
def test_x_litellm_tags_merges_with_header_tags(gateway: Gateway, tmp_path: Path, route: str) -> None:
|
||||
with wire_server(provider_reply) as wire:
|
||||
config: Final = _owned_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
|
|
@ -960,13 +889,12 @@ def test_x_litellm_tags_merges_with_header_tags(gateway: Gateway, tmp_path: Path
|
|||
]
|
||||
|
||||
|
||||
# E5: three identical requests write three spend rows, each with the tags
|
||||
@pytest.mark.parametrize(
|
||||
"route", [pytest.param("/anthropic/v1/messages", id="passthrough"), pytest.param("/v1/messages", id="unified")]
|
||||
)
|
||||
def test_repeated_requests_each_record_tags(gateway: Gateway, tmp_path: Path, route: str) -> None:
|
||||
with wire_server(provider_reply) as wire:
|
||||
config: Final = _owned_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
|
|
@ -997,10 +925,9 @@ def test_repeated_requests_each_record_tags(gateway: Gateway, tmp_path: Path, ro
|
|||
) == [EXPECTED_TAGS]
|
||||
|
||||
|
||||
# E6: guardrail mode-by-tag decider behaves identically with and without the fix
|
||||
def test_guardrail_mode_tag_decider_is_unchanged_on_pass_through(gateway: Gateway, tmp_path: Path) -> None:
|
||||
with wire_server(provider_reply) as wire:
|
||||
config: Final = _owned_config(
|
||||
config: Final = write_config(
|
||||
tmp_path,
|
||||
{
|
||||
"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]},
|
||||
|
|
@ -1048,5 +975,5 @@ def test_guardrail_mode_tag_decider_is_unchanged_on_pass_through(gateway: Gatewa
|
|||
key=key,
|
||||
headers=SENT_HEADERS,
|
||||
)
|
||||
assert control.status_code == 500, control.text
|
||||
assert control.status_code != 200, control.text
|
||||
assert len(wire.drain()) == 0, "tag-matched guardrail should have blocked before the upstream"
|
||||
|
|
|
|||
|
|
@ -1,15 +1,17 @@
|
|||
import json
|
||||
import threading
|
||||
from hashlib import sha256
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from hashlib import sha256
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import psutil
|
||||
from integration._support.client import Gateway, eventually
|
||||
from integration._support.process import owned_proxy, owned_proxy_process
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.process import owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from integration.spend._request_tag_helpers import (
|
||||
MODEL,
|
||||
OPENAI_MODEL,
|
||||
T3,
|
||||
provider_env,
|
||||
|
|
@ -17,12 +19,16 @@ from integration.spend._request_tag_helpers import (
|
|||
write_config,
|
||||
)
|
||||
|
||||
from tests.integration._support.database import read_rows
|
||||
|
||||
MODEL: Final = "claude-sonnet-4-5-20250929"
|
||||
HEADERS: Final = {"user-agent": "claude-cli/2.0.0", "x-tenant-id": "tenant-a"}
|
||||
ANTHROPIC_HEADERS: Final = {**HEADERS, "anthropic-version": "2023-06-01"}
|
||||
EXPECTED: Final = T3
|
||||
ROUTES: Final = (
|
||||
"/anthropic/v1/messages",
|
||||
"/openai/v1/chat/completions",
|
||||
"/v1/chat/completions",
|
||||
"/v1/messages",
|
||||
"/v1/responses",
|
||||
)
|
||||
|
||||
|
||||
def _ids(response) -> str:
|
||||
|
|
@ -32,77 +38,118 @@ def _ids(response) -> str:
|
|||
raise AssertionError(f"no upstream id in {response.text[:200]}")
|
||||
|
||||
|
||||
def _tagged_requests(candidate: Gateway, key: str, model: str, stream: bool, index: int) -> tuple:
|
||||
"""One call per route shape, all with the same client headers; returns (response, request_id)."""
|
||||
def _tagged_requests(
|
||||
candidate: Gateway, key: str, anthropic_model: str, openai_model: str, stream: bool, index: int
|
||||
) -> tuple:
|
||||
"""One call per route in ROUTES order with the same client headers."""
|
||||
marker: Final = f"burst {index}"
|
||||
anthropic: Final = candidate.request(
|
||||
"POST",
|
||||
"/anthropic/v1/messages",
|
||||
{"model": MODEL, "max_tokens": 16, "messages": [{"role": "user", "content": marker}], "stream": stream},
|
||||
key=key,
|
||||
headers=ANTHROPIC_HEADERS,
|
||||
return (
|
||||
candidate.request(
|
||||
"POST",
|
||||
"/anthropic/v1/messages",
|
||||
{"model": MODEL, "max_tokens": 16, "messages": [{"role": "user", "content": marker}], "stream": stream},
|
||||
key=key,
|
||||
headers=ANTHROPIC_HEADERS,
|
||||
),
|
||||
candidate.request(
|
||||
"POST",
|
||||
"/openai/v1/chat/completions",
|
||||
{"model": OPENAI_MODEL, "messages": [{"role": "user", "content": marker}], "stream": stream},
|
||||
key=key,
|
||||
headers=HEADERS,
|
||||
),
|
||||
candidate.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": openai_model, "messages": [{"role": "user", "content": marker}]},
|
||||
key=key,
|
||||
headers=HEADERS,
|
||||
),
|
||||
candidate.request(
|
||||
"POST",
|
||||
"/v1/messages",
|
||||
{
|
||||
"model": anthropic_model,
|
||||
"max_tokens": 16,
|
||||
"messages": [{"role": "user", "content": marker}],
|
||||
},
|
||||
key=key,
|
||||
headers=ANTHROPIC_HEADERS,
|
||||
),
|
||||
candidate.request("POST", "/v1/responses", {"model": openai_model, "input": marker}, key=key, headers=HEADERS),
|
||||
)
|
||||
openai: Final = candidate.request(
|
||||
"POST",
|
||||
"/openai/v1/chat/completions",
|
||||
{"model": OPENAI_MODEL, "messages": [{"role": "user", "content": marker}], "stream": stream},
|
||||
key=key,
|
||||
headers=HEADERS,
|
||||
)
|
||||
unified: Final = candidate.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "messages": [{"role": "user", "content": marker}]},
|
||||
key=key,
|
||||
headers=HEADERS,
|
||||
)
|
||||
return anthropic, openai, unified
|
||||
|
||||
|
||||
# C1: 10 concurrent bursts x 3 routes; each response id lands exactly one spend row with the tags
|
||||
def _deployments(scenario, url: str) -> tuple[str, str]:
|
||||
return (
|
||||
scenario.model(model=f"anthropic/{MODEL}", api_base=url, api_key="synthetic-anthropic-key"),
|
||||
scenario.model(model=f"openai/{OPENAI_MODEL}", api_base=f"{url}/v1"),
|
||||
)
|
||||
|
||||
|
||||
def _landed_tags(key: str, count: int) -> list[dict]:
|
||||
digest: Final = sha256(key.encode()).hexdigest()
|
||||
landed: Final = eventually(
|
||||
lambda: read_rows('SELECT request_id, request_tags FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (digest,)),
|
||||
lambda values: len(values) == count,
|
||||
seconds=70,
|
||||
)
|
||||
assert len({row["request_id"] for row in landed}) == count
|
||||
for row in landed:
|
||||
value: Final = row["request_tags"]
|
||||
assert (json.loads(value) if isinstance(value, str) else value) == EXPECTED
|
||||
return landed
|
||||
|
||||
|
||||
def _worker_pids(owned) -> tuple[int, ...]:
|
||||
workers: Final = tuple(
|
||||
child
|
||||
for child in psutil.Process(owned.process.pid).children(recursive=True)
|
||||
if any(marker in " ".join(child.cmdline()) for marker in ("spawn_main", "integration._support.proxy"))
|
||||
)
|
||||
return tuple(worker.pid for worker in workers)
|
||||
|
||||
|
||||
def test_burst_across_routes_records_tags_once_per_response(gateway: Gateway, tmp_path: Path) -> None:
|
||||
with wire_server(provider_reply) as wire:
|
||||
config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
owned_proxy_process(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as owned,
|
||||
owned.gateway.scenario() as scenario,
|
||||
):
|
||||
model: Final = scenario.model(model=f"openai/{OPENAI_MODEL}", api_base=f"{wire.url}/v1")
|
||||
candidate: Final = owned.gateway
|
||||
anthropic_model, openai_model = _deployments(scenario, wire.url)
|
||||
key: Final = scenario.key()
|
||||
|
||||
def burst(index: int) -> tuple:
|
||||
return _tagged_requests(candidate, key, model, stream=index % 2 == 1, index=index)
|
||||
return _tagged_requests(
|
||||
candidate, key, anthropic_model, openai_model, stream=index % 2 == 1, index=index
|
||||
)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=10) as pool:
|
||||
responses: Final = [response for group in pool.map(burst, range(10)) for response in group]
|
||||
assert len(responses) == 30
|
||||
assert len(responses) == 50
|
||||
assert all(response.status_code == 200 for response in responses), [
|
||||
(response.status_code, response.text[:200]) for response in responses
|
||||
]
|
||||
ids: Final = [_ids(response) for response in responses]
|
||||
assert len(set(ids)) == 30, "duplicate upstream id in burst"
|
||||
assert len(wire.drain()) == 30
|
||||
digest: Final = sha256(key.encode()).hexdigest()
|
||||
landed: Final = eventually(
|
||||
lambda: read_rows(
|
||||
'SELECT request_id, request_tags FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (digest,)
|
||||
),
|
||||
lambda values: len(values) == 30,
|
||||
seconds=70,
|
||||
)
|
||||
assert len({row["request_id"] for row in landed}) == 30
|
||||
for row in landed:
|
||||
value: Final = row["request_tags"]
|
||||
assert (json.loads(value) if isinstance(value, str) else value) == EXPECTED
|
||||
assert len(set(ids)) == 50, "duplicate upstream id in burst"
|
||||
assert len(wire.drain()) == 50
|
||||
pids: Final = _worker_pids(owned)
|
||||
assert len(set(pids)) == 2, f"expected two uvicorn workers, found {pids}"
|
||||
assert all(worker.is_running() for worker in psutil.process_iter(pids))
|
||||
_landed_tags(key, 50)
|
||||
|
||||
|
||||
# C2: generic_api sink down mid burst; spend rows still land exactly once with the tags
|
||||
def test_sink_outage_does_not_lose_spend_log_tags(gateway: Gateway, tmp_path: Path) -> None:
|
||||
stopped: Final = threading.Event()
|
||||
down: Final = threading.Event()
|
||||
delivered: Final = [] # mutable-ok: sink thread appends between drains
|
||||
|
||||
def stoppable_sink(request: Request) -> Reply:
|
||||
stopped.wait(timeout=30)
|
||||
return Reply(status=503)
|
||||
if down.is_set():
|
||||
return Reply(status=503)
|
||||
delivered.append(request)
|
||||
return Reply()
|
||||
|
||||
with wire_server(provider_reply) as wire, wire_server(stoppable_sink) as endpoint:
|
||||
config: Final = write_config(
|
||||
|
|
@ -116,7 +163,7 @@ def test_sink_outage_does_not_lose_spend_log_tags(gateway: Gateway, tmp_path: Pa
|
|||
},
|
||||
)
|
||||
with (
|
||||
owned_proxy(
|
||||
owned_proxy_process(
|
||||
gateway,
|
||||
tmp_path,
|
||||
{
|
||||
|
|
@ -126,45 +173,65 @@ def test_sink_outage_does_not_lose_spend_log_tags(gateway: Gateway, tmp_path: Pa
|
|||
},
|
||||
config=config,
|
||||
workers=2,
|
||||
) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
) as owned,
|
||||
owned.gateway.scenario() as scenario,
|
||||
):
|
||||
model: Final = scenario.model(model=f"openai/{OPENAI_MODEL}", api_base=f"{wire.url}/v1")
|
||||
candidate: Final = owned.gateway
|
||||
anthropic_model, openai_model = _deployments(scenario, wire.url)
|
||||
key: Final = scenario.key()
|
||||
|
||||
def burst(index: int) -> tuple:
|
||||
return _tagged_requests(candidate, key, model, stream=False, index=index)
|
||||
return _tagged_requests(candidate, key, anthropic_model, openai_model, stream=False, index=index)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=6) as pool:
|
||||
first: Final = [response for group in pool.map(burst, range(6)) for response in group]
|
||||
stopped.set() # sink goes down: the peer now returns 503 to every flush
|
||||
with ThreadPoolExecutor(max_workers=5) as pool:
|
||||
first: Final = [response for group in pool.map(burst, range(4)) for response in group]
|
||||
assert all(response.status_code == 200 for response in first), [
|
||||
(response.status_code, response.text[:200]) for response in first
|
||||
]
|
||||
|
||||
def second_burst(index: int) -> tuple:
|
||||
return _tagged_requests(candidate, key, model, stream=False, index=100 + index)
|
||||
def call_ids(responses: list) -> set:
|
||||
return {response.headers["x-litellm-call-id"] for response in responses}
|
||||
|
||||
with ThreadPoolExecutor(max_workers=6) as pool:
|
||||
second: Final = [response for group in pool.map(second_burst, range(6)) for response in group]
|
||||
responses: Final = [*first, *second]
|
||||
first_ids: Final = call_ids(first)
|
||||
|
||||
def events_for(ids: set) -> set:
|
||||
return {
|
||||
event["litellm_call_id"]
|
||||
for batch in delivered
|
||||
for event in json.loads(batch.body)
|
||||
if event.get("litellm_call_id") in ids
|
||||
}
|
||||
|
||||
eventually(lambda: events_for(first_ids), lambda found: found == first_ids, seconds=70)
|
||||
down.set()
|
||||
with ThreadPoolExecutor(max_workers=5) as pool:
|
||||
second: Final = [
|
||||
response for group in pool.map(lambda i: burst(100 + i), range(4)) for response in group
|
||||
]
|
||||
down.clear()
|
||||
third: Final = _tagged_requests(candidate, key, anthropic_model, openai_model, False, 200)
|
||||
responses: Final = [*first, *second, *third]
|
||||
assert all(response.status_code == 200 for response in responses), [
|
||||
(response.status_code, response.text[:200]) for response in responses
|
||||
]
|
||||
ids: Final = [_ids(response) for response in responses]
|
||||
assert len(set(ids)) == len(ids), "duplicate upstream id in burst"
|
||||
digest: Final = sha256(key.encode()).hexdigest()
|
||||
landed: Final = eventually(
|
||||
lambda: read_rows(
|
||||
'SELECT request_id, request_tags FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (digest,)
|
||||
),
|
||||
lambda values: len(values) == 36,
|
||||
second_ids: Final = call_ids(second)
|
||||
third_ids: Final = call_ids(list(third))
|
||||
recovery_probe: Final = next(iter(third_ids))
|
||||
eventually(
|
||||
lambda: events_for(third_ids),
|
||||
lambda found: recovery_probe in found,
|
||||
seconds=70,
|
||||
)
|
||||
assert len({row["request_id"] for row in landed}) == 36
|
||||
for row in landed:
|
||||
value: Final = row["request_tags"]
|
||||
assert (json.loads(value) if isinstance(value, str) else value) == EXPECTED
|
||||
second_delivered: Final = eventually(
|
||||
lambda: events_for(second_ids),
|
||||
lambda found: len(found) == len(second_ids),
|
||||
seconds=30,
|
||||
return_last_on_timeout=True,
|
||||
)
|
||||
assert second_delivered == second_ids
|
||||
_landed_tags(key, len(responses))
|
||||
|
||||
|
||||
# C3: killing one proxy worker mid burst loses no spend row
|
||||
def test_worker_kill_mid_burst_loses_no_spend_rows(gateway: Gateway, tmp_path: Path) -> None:
|
||||
with wire_server(provider_reply) as wire:
|
||||
config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
|
||||
|
|
@ -173,41 +240,38 @@ def test_worker_kill_mid_burst_loses_no_spend_rows(gateway: Gateway, tmp_path: P
|
|||
owned.gateway.scenario() as scenario,
|
||||
):
|
||||
candidate: Final = owned.gateway
|
||||
model: Final = scenario.model(model=f"openai/{OPENAI_MODEL}", api_base=f"{wire.url}/v1")
|
||||
anthropic_model, openai_model = _deployments(scenario, wire.url)
|
||||
key: Final = scenario.key()
|
||||
|
||||
def burst(index: int) -> tuple:
|
||||
return _tagged_requests(candidate, key, model, stream=False, index=index)
|
||||
return _tagged_requests(candidate, key, anthropic_model, openai_model, stream=False, index=index)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=6) as pool:
|
||||
first: Final = [response for group in pool.map(burst, range(6)) for response in group]
|
||||
with ThreadPoolExecutor(max_workers=5) as pool:
|
||||
first: Final = [response for group in pool.map(burst, range(4)) for response in group]
|
||||
|
||||
workers: Final = [
|
||||
child
|
||||
for child in psutil.Process(owned.process.pid).children(recursive=True)
|
||||
if child.status() != psutil.STATUS_ZOMBIE
|
||||
if any(marker in " ".join(child.cmdline()) for marker in ("spawn_main", "integration._support.proxy"))
|
||||
]
|
||||
assert len(workers) >= 2, f"expected two proxy workers, found {[w.pid for w in workers]}"
|
||||
assert len(workers) == 2, (
|
||||
f"expected two uvicorn workers, found {[(w.pid, w.cmdline()[:3]) for w in workers]}"
|
||||
)
|
||||
workers[0].kill()
|
||||
psutil.wait_procs(workers[:1], timeout=10)
|
||||
assert not workers[0].is_running()
|
||||
|
||||
def second_burst(index: int) -> tuple:
|
||||
return _tagged_requests(candidate, key, model, stream=False, index=100 + index)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=6) as pool:
|
||||
second: Final = [response for group in pool.map(second_burst, range(6)) for response in group]
|
||||
with ThreadPoolExecutor(max_workers=5) as pool:
|
||||
second: Final = [
|
||||
response for group in pool.map(lambda i: burst(100 + i), range(4)) for response in group
|
||||
]
|
||||
responses: Final = [*first, *second]
|
||||
for position in range(len(ROUTES)):
|
||||
statuses: Final = {
|
||||
responses[offset + position].status_code for offset in range(0, len(responses), len(ROUTES))
|
||||
}
|
||||
assert 200 in statuses, f"no surviving 200 for route {ROUTES[position]}: {statuses}"
|
||||
ok: Final = [response for response in responses if response.status_code == 200]
|
||||
ids: Final = [_ids(response) for response in ok]
|
||||
assert len(set(ids)) == len(ids), "duplicate upstream id in burst"
|
||||
digest: Final = sha256(key.encode()).hexdigest()
|
||||
landed: Final = eventually(
|
||||
lambda: read_rows(
|
||||
'SELECT request_id, request_tags FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (digest,)
|
||||
),
|
||||
lambda values: len(values) == len(ids),
|
||||
seconds=70,
|
||||
)
|
||||
assert len({row["request_id"] for row in landed}) == len(ids)
|
||||
for row in landed:
|
||||
value: Final = row["request_tags"]
|
||||
assert (json.loads(value) if isinstance(value, str) else value) == EXPECTED
|
||||
_landed_tags(key, len(ok))
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue