diff --git a/tests/integration/spend/_request_tag_helpers.py b/tests/integration/spend/_request_tag_helpers.py index c7871862c11..4cf6b16d475 100644 --- a/tests/integration/spend/_request_tag_helpers.py +++ b/tests/integration/spend/_request_tag_helpers.py @@ -116,6 +116,7 @@ def provider_reply(request: Request) -> Reply: """Scripted edge for every route the audit drives: anthropic messages, openai chat completions and responses, gemini generateContent. Error bodies keyed off the sentinel model name.""" body: Final = json.loads(request.body) if request.body else {} + target: Final = request.target.split("?", 1)[0] if body.get("model") == "claude-nonexistent-model": return Reply( status=400, @@ -123,9 +124,9 @@ def provider_reply(request: Request) -> Reply: {"error": {"type": "invalid_request_error", "message": "model: claude-nonexistent-model"}} ).encode(), ) - if request.target == "/v1/models" or request.target.startswith("/v1/models/"): + if target == "/v1/models" or target.startswith("/v1/models/"): return Reply(body=json.dumps({"object": "list", "data": []}).encode()) - if request.target == "/v1/messages": + if target == "/v1/messages": identity: Final = _message_id() if body.get("stream") is True: events: Final = tuple( @@ -134,7 +135,7 @@ def provider_reply(request: Request) -> Reply: ) return Reply(content_type="text/event-stream", chunks=_sse_frames(events)) return Reply(body=json.dumps({**ANTHROPIC_SONNET_BODY, "id": identity}).encode()) - if request.target == "/v1/chat/completions": + if target == "/v1/chat/completions": identity = "chatcmpl_" + uuid.uuid4().hex if body.get("stream") is True: frames: Final = tuple( @@ -142,7 +143,7 @@ def provider_reply(request: Request) -> Reply: ) + (b"data: [DONE]\n\n",) return Reply(content_type="text/event-stream", chunks=frames) return Reply(body=json.dumps({**CHAT_COMPLETION_BODY, "id": identity}).encode()) - if request.target == "/v1/responses": + if target == "/v1/responses": identity = "resp_" + uuid.uuid4().hex message: Final = _message_id() completed_body: Final = { @@ -165,7 +166,7 @@ def provider_reply(request: Request) -> Reply: completed: Final = {"type": "response.completed", "response": completed_body} return Reply(content_type="text/event-stream", chunks=_sse_frames((created, delta, completed))) return Reply(body=json.dumps(completed_body).encode()) - if request.target.endswith(":generateContent") or request.target.endswith(":streamGenerateContent"): + if target.endswith(":generateContent") or target.endswith(":streamGenerateContent"): return Reply(body=json.dumps(GEMINI_BODY).encode()) raise AssertionError(f"unexpected upstream target {request.target}") diff --git a/tests/integration/spend/test_spend_log_request_tags.py b/tests/integration/spend/test_spend_log_request_tags.py index cb57496e52e..89b8c4609c3 100644 --- a/tests/integration/spend/test_spend_log_request_tags.py +++ b/tests/integration/spend/test_spend_log_request_tags.py @@ -1,4 +1,5 @@ import asyncio +import http.client import json import uuid from hashlib import sha256 @@ -19,6 +20,7 @@ from integration.spend._request_tag_helpers import ( provider_reply, tags_by_id, tags_by_key, + tags_of, write_config, ) @@ -555,22 +557,30 @@ def test_routes_without_headers_record_no_tags(gateway: Gateway, tmp_path: Path, model=f"anthropic/{MODEL}", api_base=wire.url, api_key="synthetic-anthropic-key" ) key: Final = scenario.key() - response: Final = candidate.request( + connection: Final = http.client.HTTPConnection("127.0.0.1", candidate.client.base_url.port) + connection.request( "POST", route, - { - "model": MODEL if route == "/anthropic/v1/messages" else model, - "max_tokens": 16, - "messages": [{"role": "user", "content": f"tag me {uuid.uuid4().hex}"}], + body=json.dumps( + { + "model": MODEL if route == "/anthropic/v1/messages" else model, + "max_tokens": 16, + "messages": [{"role": "user", "content": f"tag me {uuid.uuid4().hex}"}], + } + ), + headers={ + "authorization": f"Bearer {key}", + "content-type": "application/json", + "anthropic-version": "2023-06-01", }, - key=key, - headers={"user-agent": "", "anthropic-version": "2023-06-01"}, ) - assert response.status_code == 200, response.text + raw: Final = connection.getresponse() + payload: Final = raw.read() + connection.close() + assert raw.status == 200, payload + request_id: Final = json.loads(payload)["id"] assert len(wire.drain()) == 1 - assert eventually(lambda: tags_by_id(response.json()["id"]), lambda tags: len(tags) == 1, seconds=70) == [ - [] - ] + 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 @@ -701,6 +711,15 @@ def test_unauthenticated_request_writes_no_spend_row(gateway: Gateway, tmp_path: headers={**SENT_HEADERS, "anthropic-version": "2023-06-01"}, ) assert response.status_code == 401, response.text + anonymous: Final = eventually( + lambda: read_rows( + "SELECT request_tags FROM \"LiteLLM_SpendLogs\" WHERE api_key IS NULL OR api_key=''", + (), + ), + lambda rows: len(rows) == 1, + seconds=70, + ) + assert [tags_of(row) for row in anonymous] == [[]] key: Final = scenario.key() control: Final = candidate.request( "POST", @@ -718,7 +737,7 @@ 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 + 1 + assert _spend_count() == before + 2 # S7: an upstream 400 surfaces the same status and its spend row records the tags