mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-05 02:41:56 +00:00
41df8cf4d0
* wip * feat(traces): establish shared Rust storage foundation * fix(traces): escape ClickHouse text parameters * test(traces): exercise response cap with bounded strings * fix(traces): remove unnecessary lint expectation * fix(traces): encode ClickHouse timestamp units in Rust * test(traces): mark exception match as a regex * refactor(traces): execute schema setup in Rust * refactor(traces): use shared logging execution wrapper * docs(traces): replace foundation README with boundary rules * fix(traces): use current bridge execution facade * fix(traces): account for protocol cast in lint budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
78 lines
3.4 KiB
Python
78 lines
3.4 KiB
Python
import base64
|
|
import json
|
|
from typing import Final
|
|
from urllib.parse import parse_qs, urlsplit
|
|
|
|
import pytest
|
|
|
|
from litellm.rust_bridge.traces import encode_rows, ensure_schema, query
|
|
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
|
|
|
|
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"}]}))
|
|
rows: Final = await query(
|
|
recording_server.base_url + "?database=wrong&user=wrong&password=wrong",
|
|
"trace_test",
|
|
"reader",
|
|
"p@ss/word%",
|
|
"SELECT {trace_id:String} AS trace_id",
|
|
{"trace_id": "trace-1"},
|
|
)
|
|
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"
|
|
assert parameters["database"] == ["trace_test"]
|
|
assert parameters["param_trace_id"] == ["trace-1"]
|
|
assert parameters["readonly"] == ["1"]
|
|
assert "user" not in parameters
|
|
assert "password" not in parameters
|
|
assert request.headers["authorization"] == "Basic " + base64.b64encode(b"reader:p@ss/word%").decode()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_trace_reader_rejects_success_status_with_embedded_error(recording_server: RecordingServer) -> None:
|
|
recording_server.enqueue(ResponseSpec(body={"data": [], "exception": "query failed"}))
|
|
with pytest.raises(RuntimeError, match="invalid or failed JSON"):
|
|
await query(recording_server.base_url, "trace_test", "reader", "password", "SELECT 1", {})
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("database,retention", [("db; DROP DATABASE default", 7), ("traces", 0)])
|
|
async def test_schema_binding_preserves_configuration_validation(database: str, retention: int) -> None:
|
|
with pytest.raises(ValueError, match=r"database.*retention"):
|
|
await ensure_schema("http://localhost:8123", database, "writer", "password", retention, 14)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schema_setup_uses_writer_credentials_and_rejects_failed_statement(recording_server: RecordingServer) -> None:
|
|
recording_server.expected_requests = 2
|
|
recording_server.enqueue(ResponseSpec(body=""))
|
|
recording_server.enqueue(ResponseSpec(status=403, body="denied"))
|
|
with pytest.raises(RuntimeError, match="schema setup failed with HTTP status 403"):
|
|
await ensure_schema(
|
|
recording_server.base_url + "?database=wrong&readonly=1",
|
|
"trace_test",
|
|
"writer",
|
|
"p@ss/word%",
|
|
7,
|
|
14,
|
|
)
|
|
assert len(recording_server.requests) == 2
|
|
assert recording_server.requests[0].raw_body.startswith(b"CREATE DATABASE IF NOT EXISTS")
|
|
assert recording_server.requests[1].raw_body.startswith(b"CREATE TABLE IF NOT EXISTS")
|
|
assert "readonly" not in parse_qs(urlsplit(recording_server.requests[0].path).query)
|
|
assert recording_server.requests[0].headers["authorization"] == "Basic " + base64.b64encode(
|
|
b"writer:p@ss/word%"
|
|
).decode()
|
|
|
|
|
|
def test_insert_encoding_preserves_nanoseconds_through_bridge() -> None:
|
|
assert json.loads(encode_rows([{"Timestamp": 1_234_567_890, "Input": "hello"}])) == {
|
|
"Input": "hello",
|
|
"Timestamp": "1970-01-01T00:00:01.23456789Z",
|
|
}
|