76 lines
2.0 KiB
Python
76 lines
2.0 KiB
Python
|
"""Regression tests for PostgreSQL audit persistence."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
|
||
|
|
BACKEND_DIR = Path(__file__).resolve().parents[1]
|
||
|
|
if str(BACKEND_DIR) not in __import__("sys").path:
|
||
|
|
__import__("sys").path.insert(0, str(BACKEND_DIR))
|
||
|
|
|
||
|
|
import audit_store # noqa: E402
|
||
|
|
from audit_store import PostgresAuditStore # noqa: E402
|
||
|
|
|
||
|
|
|
||
|
|
class _RecordingCursor:
|
||
|
|
def __init__(self) -> None:
|
||
|
|
self.query = ""
|
||
|
|
self.params = ()
|
||
|
|
|
||
|
|
def __enter__(self):
|
||
|
|
return self
|
||
|
|
|
||
|
|
def __exit__(self, exc_type, exc, tb):
|
||
|
|
return False
|
||
|
|
|
||
|
|
def execute(self, query: str, params=()) -> None:
|
||
|
|
self.query = query
|
||
|
|
self.params = params or ()
|
||
|
|
assert query.count("%s") == len(self.params)
|
||
|
|
|
||
|
|
|
||
|
|
class _RecordingConnection:
|
||
|
|
def __init__(self, cursor: _RecordingCursor) -> None:
|
||
|
|
self._cursor = cursor
|
||
|
|
|
||
|
|
def __enter__(self):
|
||
|
|
return self
|
||
|
|
|
||
|
|
def __exit__(self, exc_type, exc, tb):
|
||
|
|
return False
|
||
|
|
|
||
|
|
def cursor(self) -> _RecordingCursor:
|
||
|
|
return self._cursor
|
||
|
|
|
||
|
|
|
||
|
|
def test_record_llm_call_keeps_columns_placeholders_and_params_aligned(monkeypatch):
|
||
|
|
cursor = _RecordingCursor()
|
||
|
|
monkeypatch.setattr(audit_store, "psycopg", object())
|
||
|
|
store = PostgresAuditStore("postgresql://unused")
|
||
|
|
store._initialized = True
|
||
|
|
monkeypatch.setattr(store, "_connect", lambda: _RecordingConnection(cursor))
|
||
|
|
|
||
|
|
store.record_llm_call({
|
||
|
|
"request_id": "request-1",
|
||
|
|
"session_hash": "session",
|
||
|
|
"ip_hash": "ip",
|
||
|
|
"job_type": "ocr",
|
||
|
|
"model": "vision-model",
|
||
|
|
"estimated_input_tokens": 12,
|
||
|
|
"max_output_tokens": 256,
|
||
|
|
"estimated_cost": 0.01,
|
||
|
|
"actual_output_chars": 42,
|
||
|
|
"actual_cost": 0.02,
|
||
|
|
"queue_ms": 10,
|
||
|
|
"run_ms": 20,
|
||
|
|
"total_ms": 30,
|
||
|
|
"status": "completed",
|
||
|
|
"error_code": "",
|
||
|
|
"metadata": {"source": "test"},
|
||
|
|
})
|
||
|
|
|
||
|
|
assert "INSERT INTO llm_call_audit" in cursor.query
|
||
|
|
assert cursor.query.count("%s") == 16
|
||
|
|
assert len(cursor.params) == 16
|