160 lines
5.0 KiB
Python
160 lines
5.0 KiB
Python
"""PostgresRecorder 集成测试(M2 设计 §5;真实实验室 Postgres,polygateway 专用库)。
|
|
|
|
DSN 走 .env `PGW_TELEMETRY_PG_DSN`,缺则 skip。该实例上有 app/chs_prod 等
|
|
在用库——本测试只允许连 polygateway 专用库(fixture 里守卫)。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import os
|
|
|
|
import pytest
|
|
from dotenv import dotenv_values
|
|
|
|
from polygateway.telemetry.postgres import PostgresRecorder
|
|
|
|
_EXPECTED_COLUMNS = [
|
|
"call_id",
|
|
"parent_call_id",
|
|
"session_id",
|
|
"model",
|
|
"provider",
|
|
"source_name",
|
|
"messages",
|
|
"response",
|
|
"thinking",
|
|
"prompt_tokens",
|
|
"completion_tokens",
|
|
"usage_source",
|
|
"latency_ms",
|
|
"ttft_ms",
|
|
"max_inter_token_ms",
|
|
"cache_hit",
|
|
"error",
|
|
"cost",
|
|
"created_at",
|
|
]
|
|
|
|
|
|
def _dsn() -> str | None:
|
|
merged = {**dotenv_values(".env"), **os.environ}
|
|
raw = merged.get("PGW_TELEMETRY_PG_DSN")
|
|
if not raw:
|
|
return None
|
|
scheme, sep, rest = raw.partition("://")
|
|
return f"{scheme.partition('+')[0]}{sep}{rest}"
|
|
|
|
|
|
@pytest.fixture
|
|
async def dsn():
|
|
value = _dsn()
|
|
if value is None:
|
|
pytest.skip("PGW_TELEMETRY_PG_DSN 未配置")
|
|
# 隔离守卫: 该实例有 app/chs_prod/mimiciv 等在用库,只许打 polygateway 专用库
|
|
if not value.rstrip("/").endswith("/polygateway"):
|
|
pytest.fail(f"遥测测试只允许连 polygateway 专用库,当前 DSN 库名不符: {value!r}")
|
|
import asyncpg
|
|
|
|
conn = await asyncpg.connect(value, timeout=10)
|
|
try:
|
|
await conn.execute("DROP TABLE IF EXISTS llm_calls")
|
|
finally:
|
|
await conn.close()
|
|
return value
|
|
|
|
|
|
async def _record_minimal(recorder: PostgresRecorder, call_id: str = "c1", **overrides) -> None:
|
|
fields = {
|
|
"call_id": call_id,
|
|
"parent_call_id": None,
|
|
"session_id": "sess-1",
|
|
"model": "m",
|
|
"provider": "p",
|
|
"source_name": "s1",
|
|
"messages": "[]",
|
|
"response": "ok",
|
|
"thinking": "",
|
|
"prompt_tokens": 1,
|
|
"completion_tokens": 2,
|
|
"usage_source": "measured",
|
|
"latency_ms": 10,
|
|
"ttft_ms": None,
|
|
"max_inter_token_ms": None,
|
|
"cache_hit": False,
|
|
"error": None,
|
|
"cost": None,
|
|
}
|
|
fields.update(overrides)
|
|
await recorder.record_llm_call(**fields)
|
|
|
|
|
|
async def _fetch(dsn: str, sql: str):
|
|
import asyncpg
|
|
|
|
conn = await asyncpg.connect(dsn, timeout=10)
|
|
try:
|
|
return await conn.fetch(sql)
|
|
finally:
|
|
await conn.close()
|
|
|
|
|
|
class TestSchema:
|
|
async def test_schema_has_frozen_columns_in_order(self, dsn):
|
|
recorder = PostgresRecorder(dsn)
|
|
try:
|
|
await _record_minimal(recorder)
|
|
rows = await _fetch(
|
|
dsn,
|
|
"SELECT column_name FROM information_schema.columns "
|
|
"WHERE table_name='llm_calls' ORDER BY ordinal_position",
|
|
)
|
|
assert [r["column_name"] for r in rows] == _EXPECTED_COLUMNS
|
|
finally:
|
|
await recorder.aclose()
|
|
|
|
async def test_call_id_idempotent(self, dsn):
|
|
recorder = PostgresRecorder(dsn)
|
|
try:
|
|
await _record_minimal(recorder, call_id="dup")
|
|
await _record_minimal(recorder, call_id="dup", response="second")
|
|
rows = await _fetch(dsn, "SELECT response FROM llm_calls WHERE call_id='dup'")
|
|
assert [r["response"] for r in rows] == ["ok"] # ON CONFLICT DO NOTHING
|
|
finally:
|
|
await recorder.aclose()
|
|
|
|
async def test_concurrent_writes_all_land(self, dsn):
|
|
recorder = PostgresRecorder(dsn)
|
|
try:
|
|
await asyncio.gather(*(_record_minimal(recorder, call_id=f"c{i}") for i in range(50)))
|
|
rows = await _fetch(dsn, "SELECT count(*) AS n FROM llm_calls")
|
|
assert rows[0]["n"] == 50
|
|
finally:
|
|
await recorder.aclose()
|
|
|
|
|
|
class TestDegradation:
|
|
async def test_unreachable_server_degrades_silently(self):
|
|
"""结构性失败(建池不通)→ warning 一次后永久降级,业务零感知。"""
|
|
recorder = PostgresRecorder("postgresql://u:p@127.0.0.1:1/x")
|
|
await _record_minimal(recorder) # 不抛
|
|
await _record_minimal(recorder, call_id="c2") # 已降级短路,同样不抛
|
|
await recorder.aclose()
|
|
|
|
async def test_row_failure_does_not_poison_later_rows(self, dsn):
|
|
"""运行时单条写失败(NUL 字节文本被 PG 拒)→ 丢该行,后续行照常落库。"""
|
|
recorder = PostgresRecorder(dsn)
|
|
try:
|
|
await _record_minimal(recorder, call_id="bad", response="nul\x00byte")
|
|
await _record_minimal(recorder, call_id="good")
|
|
rows = await _fetch(dsn, "SELECT call_id FROM llm_calls ORDER BY call_id")
|
|
assert [r["call_id"] for r in rows] == ["good"]
|
|
finally:
|
|
await recorder.aclose()
|
|
|
|
async def test_aclose_idempotent(self, dsn):
|
|
recorder = PostgresRecorder(dsn)
|
|
await _record_minimal(recorder)
|
|
await recorder.aclose()
|
|
await recorder.aclose()
|