"""PostgresRecorder 集成测试(M2 设计 §5;真实实验室 Postgres,polygateway 专用库)。 DSN 走 .env `PGW_TELEMETRY_PG_DSN`,缺则 skip。该实例上有 app/chs_prod 等 在用库——本测试只允许连 polygateway 专用库(fixture 里守卫)。 隔离纪律(M4 事故教训): `llm_calls` 是与真实批跑/迁移项目共享的表, **严禁 DROP/TRUNCATE**——本测试以 run 级 call_id 前缀隔离,断言只看 自己写入的行,teardown 只删自己的行。 """ from __future__ import annotations import asyncio import os from uuid import uuid4 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", "cached_prompt_tokens", "model_reported", ] # run 级前缀: 同库并存的其他运行(迁移批跑/另一开发机)互不可见 _RUN_PREFIX = f"pgwtest-{uuid4().hex[:8]}" def _cid(suffix: str) -> str: return f"{_RUN_PREFIX}-{suffix}" 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}") yield value # teardown: 只删本 run 写入的行;表可能尚不存在(全新库)则忽略 import asyncpg conn = await asyncpg.connect(value, timeout=10) try: if await conn.fetchval("SELECT to_regclass('llm_calls')") is not None: await conn.execute("DELETE FROM llm_calls WHERE call_id LIKE $1", f"{_RUN_PREFIX}-%") finally: await conn.close() async def _record_minimal( recorder: PostgresRecorder, call_id: str | None = None, **overrides ) -> None: fields = { "call_id": call_id if call_id is not None else _cid("c1"), "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, "cached_prompt_tokens": None, "model_reported": None, } fields.update(overrides) await recorder.record_llm_call(**fields) async def _fetch(dsn: str, sql: str, *args): import asyncpg conn = await asyncpg.connect(dsn, timeout=10) try: return await conn.fetch(sql, *args) finally: await conn.close() _LEGACY_DDL = """ CREATE TABLE {schema}.llm_calls ( call_id TEXT PRIMARY KEY, parent_call_id TEXT, session_id TEXT, model TEXT NOT NULL, provider TEXT NOT NULL, source_name TEXT NOT NULL, messages TEXT NOT NULL, response TEXT NOT NULL, thinking TEXT NOT NULL DEFAULT '', prompt_tokens INTEGER NOT NULL, completion_tokens INTEGER NOT NULL, usage_source TEXT NOT NULL, latency_ms INTEGER NOT NULL, ttft_ms DOUBLE PRECISION, max_inter_token_ms DOUBLE PRECISION, cache_hit BOOLEAN NOT NULL DEFAULT FALSE, error TEXT, cost DOUBLE PRECISION, created_at TIMESTAMPTZ NOT NULL DEFAULT now() ) """ @pytest.fixture async def legacy_schema(dsn): """在**自建的临时 schema** 里造一张 18 列旧表,验证补列(issue #3)。 绝不碰共享的 public.llm_calls: 用 search_path 把 recorder 指向临时 schema, teardown 只 DROP 自己建的 schema。 """ import asyncpg name = f"pgwtest_{uuid4().hex[:8]}" conn = await asyncpg.connect(dsn, timeout=10) try: await conn.execute(f"CREATE SCHEMA {name}") await conn.execute(_LEGACY_DDL.format(schema=name)) finally: await conn.close() sep = "&" if "?" in dsn else "?" yield f"{dsn}{sep}options=-csearch_path%3D{name}", name conn = await asyncpg.connect(dsn, timeout=10) try: await conn.execute(f"DROP SCHEMA {name} CASCADE") finally: await conn.close() class TestObservabilityColumns: """issue #3: 两列写入可回读,且已存在的 18 列旧表会被自动补列。""" async def test_values_round_trip(self, dsn): recorder = PostgresRecorder(dsn) try: await _record_minimal(recorder, call_id=_cid("hit"), cached_prompt_tokens=64) await _record_minimal(recorder, call_id=_cid("zero"), cached_prompt_tokens=0) await _record_minimal(recorder, call_id=_cid("model"), model_reported="MiniMax-01") rows = await _fetch( dsn, "SELECT call_id, cached_prompt_tokens, model_reported FROM llm_calls " "WHERE call_id LIKE $1", f"{_RUN_PREFIX}-%", ) by_id = {r["call_id"]: r for r in rows} assert by_id[_cid("hit")]["cached_prompt_tokens"] == 64 assert by_id[_cid("zero")]["cached_prompt_tokens"] == 0 # 真实零命中 ≠ NULL assert by_id[_cid("model")]["cached_prompt_tokens"] is None assert by_id[_cid("model")]["model_reported"] == "MiniMax-01" finally: await recorder.aclose() async def test_legacy_table_is_upgraded_in_place(self, legacy_schema): """18 列旧表不补列的话,每行写入都会被逐行 warning 丢弃(遥测静默全失)。""" schema_dsn, schema = legacy_schema recorder = PostgresRecorder(schema_dsn) try: await _record_minimal( recorder, call_id=_cid("legacy"), cached_prompt_tokens=7, model_reported="m-real" ) cols = await _fetch( schema_dsn, "SELECT column_name FROM information_schema.columns " "WHERE table_schema = $1 AND table_name = 'llm_calls' ORDER BY ordinal_position", schema, ) # ALTER 只能追加到末尾: 与新建库的列序一致才不会分叉 assert [r["column_name"] for r in cols] == _EXPECTED_COLUMNS rows = await _fetch( schema_dsn, "SELECT cached_prompt_tokens, model_reported FROM llm_calls WHERE call_id = $1", _cid("legacy"), ) assert (rows[0]["cached_prompt_tokens"], rows[0]["model_reported"]) == (7, "m-real") finally: await recorder.aclose() 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=_cid("dup")) await _record_minimal(recorder, call_id=_cid("dup"), response="second") rows = await _fetch( dsn, "SELECT response FROM llm_calls WHERE call_id = $1", _cid("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=_cid(f"c{i}")) for i in range(50)) ) rows = await _fetch( dsn, "SELECT count(*) AS n FROM llm_calls WHERE call_id LIKE $1", f"{_RUN_PREFIX}-c%", ) 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=_cid("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=_cid("bad"), response="nul\x00byte") await _record_minimal(recorder, call_id=_cid("good")) rows = await _fetch( dsn, "SELECT call_id FROM llm_calls WHERE call_id = ANY($1::text[]) ORDER BY call_id", [_cid("bad"), _cid("good")], ) assert [r["call_id"] for r in rows] == [_cid("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()