diff --git a/tests/test_litellm_rust/test_traces.py b/tests/test_litellm_rust/test_traces.py index 447ca2ce4bb..5f628a96edf 100644 --- a/tests/test_litellm_rust/test_traces.py +++ b/tests/test_litellm_rust/test_traces.py @@ -1,6 +1,7 @@ import base64 import gzip import json +import time from typing import Final from urllib.parse import parse_qs, urlsplit @@ -14,16 +15,24 @@ pytestmark = pytest.mark.requires_rust_extension @pytest.mark.asyncio async def test_trace_reader_projects_connection_and_parameters(recording_server: RecordingServer) -> None: - recording_server.enqueue(ResponseSpec(body={"data": [{"trace_id": "trace-1"}]})) + recording_server.enqueue(ResponseSpec(body={"data": [{"span_id": "span-1"}]})) reader_url: Final = recording_server.base_url.replace("http://", "http://reader:p%40ss%2Fword%25@") storage: Final = NativeTraceStorage("trace_test", recording_server.base_url, reader_url + "?database=wrong") - rows: Final = json.loads(await storage.query("SELECT {trace_id:String} AS trace_id", {"trace_id": "trace-1"})) + rows: Final = json.loads( + await storage.query( + "trace_spans", {"trace_id": "trace-1", "team_ids": [], "api_key_hash": "", "trace_ref": ""} + ) + ) request: Final = recording_server.requests[0] - parameters: Final = parse_qs(urlsplit(request.path).query) - assert rows == [{"trace_id": "trace-1"}] - assert request.raw_body == b"SELECT {trace_id:String} AS trace_id" + parameters: Final = parse_qs(urlsplit(request.path).query, keep_blank_values=True) + assert rows == {"data": [{"span_id": "span-1"}]} + assert b"FROM otel_traces AS o" in request.raw_body + assert b"WHERE o.TraceId = {trace_id:String}" in request.raw_body assert parameters["database"] == ["trace_test"] assert parameters["param_trace_id"] == ["trace-1"] + assert parameters["param_team_ids"] == ["[]"] + assert parameters["param_api_key_hash"] == [""] + assert parameters["param_trace_ref"] == [""] assert parameters["readonly"] == ["1"] assert "user" not in parameters assert "password" not in parameters @@ -35,6 +44,16 @@ async def test_trace_reader_rejects_success_status_with_embedded_error(recording recording_server.enqueue(ResponseSpec(body={"data": [], "exception": "query failed"})) storage: Final = NativeTraceStorage("trace_test", recording_server.base_url, recording_server.base_url) with pytest.raises(RuntimeError, match="invalid or failed JSON"): + await storage.query( + "trace_spans", {"trace_id": "trace-1", "team_ids": [], "api_key_hash": "", "trace_ref": ""} + ) + + +@pytest.mark.asyncio +async def test_trace_reader_rejects_arbitrary_sql_before_sending(recording_server: RecordingServer) -> None: + recording_server.expected_requests = 0 + storage: Final = NativeTraceStorage("trace_test", recording_server.base_url, recording_server.base_url) + with pytest.raises(ValueError, match="unknown ClickHouse read query"): await storage.query("SELECT 1", {}) @@ -73,9 +92,15 @@ async def test_schema_setup_uses_writer_credentials_and_rejects_failed_statement async def test_insert_encodes_and_sends_rows(recording_server: RecordingServer) -> None: recording_server.enqueue(ResponseSpec(body="")) storage: Final = NativeTraceStorage("trace_test", recording_server.base_url) - await storage.insert_rows("otel_traces", [{"Timestamp": 1_234_567_890, "Input": "hello"}]) + before_insert_ms: Final = time.time_ns() // 1_000_000 + await storage.insert_rows("otel_traces", [{"Timestamp": 1_234_567_890, "Input": "hello", "EngineReceivedMs": 0}]) + after_insert_ms: Final = time.time_ns() // 1_000_000 request: Final = recording_server.requests[0] - assert json.loads(gzip.decompress(request.raw_body)) == { + row: Final = json.loads(gzip.decompress(request.raw_body)) + assert type(row["EngineReceivedMs"]) is int + assert before_insert_ms <= row["EngineReceivedMs"] <= after_insert_ms + assert row == { + "EngineReceivedMs": row["EngineReceivedMs"], "Input": "hello", "Timestamp": "1970-01-01T00:00:01.23456789Z", }