test: build the sandbox factory the PG tests will run inside
Seven copies of "create a schema, hang it off search_path, drop it in teardown" were spread across two files, each with its own cleanup. Any one of them written wrong leaves the residue on a database shared with real batch runs. This is one implementation, and it makes "the test cannot reach the admin connection" a structural fact rather than a note in a docstring. Three role modes cover every fixture that exists today: none for plain schema isolation, owner for the retention script's own runs, grantee for the least-privilege deployment cases. Owner runs its DDL as itself so it ends up owning the table; grantee is the opposite, since that case only means anything when someone else built it. The schema and the role deliberately get different prefixes. Give them the same name and "$user" resolves to the sandbox, which hides the shared table and quietly turns the worst-case test into a test of nothing. Writing it also turned up a bug in my first version: rolling back a failed sandbox unwound the whole stack, so an earlier sandbox in the same test lost its role mid-use. The test for it fails with a password authentication error, which is what that looks like from the outside. Each call now unwinds only what it created, and cleanup tries every statement before raising, since one failure stranding the rest means global roles left behind by hand.
This commit is contained in:
@@ -0,0 +1,211 @@
|
||||
"""PG 集成测试的一次性沙箱工厂(issue #18)。
|
||||
|
||||
**为什么把它收敛成一份**: 在此之前,"建临时 schema → 挂 search_path → teardown
|
||||
删净"这套样板在两个测试文件里重复了七处,清理逻辑各写各的——任何一处写漏,残留都
|
||||
落在与真实批跑共用的那个库上。工厂让清理只有一份实现,并让"用例拿不到管理连接"
|
||||
成为结构事实而不是纪律。
|
||||
|
||||
**admin DSN 不做成 fixture**: 它能对共享表执行任何语句。做成 fixture 等于把这个
|
||||
能力摆在每一条用例面前,"用例不该直接用"就只是一句提醒。故它是模块私有函数,
|
||||
只被工厂内部调用,`PgSandbox` 也不携带它。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Literal
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from dotenv import dotenv_values
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
# 测试专用口令: 这些角色只在单条用例的生命周期内存在,且只对自建 schema 有权。
|
||||
# 它不是机密,写死在这里比走 .env 更清楚——.env 里的每一项都该是真实部署会用的。
|
||||
_SANDBOX_PASSWORD = "pgw-sandbox-not-a-secret" # noqa: S105
|
||||
|
||||
_Role = Literal["none", "owner", "grantee"]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PgSandbox:
|
||||
"""一次性 PG 沙箱: 独立 schema + 可选独占登录角色。"""
|
||||
|
||||
schema: str
|
||||
role: str | None
|
||||
dsn: str
|
||||
"""已挂 `options=-csearch_path=<schema>`,用例默认用它。"""
|
||||
bare_dsn: str | None
|
||||
"""同角色但**不挂** search_path(回落 `"$user", public`);`role="none"` 时为 None。"""
|
||||
|
||||
|
||||
def _admin_dsn() -> str | None:
|
||||
"""读 `.env` 的 `PGW_TELEMETRY_PG_DSN` 并剥掉 SQLAlchemy 风格的 `+driver` 后缀。"""
|
||||
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}"
|
||||
|
||||
|
||||
def _require_admin_dsn() -> str:
|
||||
"""取管理连接串;未配置则 skip,连错库则 fail(不是 skip)。
|
||||
|
||||
库名守卫不肯降级成 skip: 这个实例上还有 app/chs_prod 等在用库,把"连错库"
|
||||
悄悄跳过,等于让一次配置事故以"没跑那些测试"的形态过关。
|
||||
"""
|
||||
value = _admin_dsn()
|
||||
if value is None:
|
||||
pytest.skip("PGW_TELEMETRY_PG_DSN 未配置")
|
||||
if not value.rstrip("/").endswith("/polygateway"):
|
||||
pytest.fail(f"PG 集成测试只允许连 polygateway 专用库,当前 DSN 库名不符: {value!r}")
|
||||
return value
|
||||
|
||||
|
||||
def _with_search_path(dsn: str, schema: str) -> str:
|
||||
sep = "&" if "?" in dsn else "?"
|
||||
return f"{dsn}{sep}options=-csearch_path%3D{schema}"
|
||||
|
||||
|
||||
def _as_role(dsn: str, role: str) -> str:
|
||||
"""把 DSN 的用户名口令段换成沙箱角色的,其余(主机/库/参数)原样保留。"""
|
||||
return re.sub(r"//[^@/]+@", f"//{role}:{_SANDBOX_PASSWORD}@", dsn, count=1)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def pg_catalog_probe():
|
||||
"""只读地查 PG catalog,**仅供工厂自测核对残留**,不是通用查询入口。
|
||||
|
||||
它拿的是管理连接,故有意只暴露给 `test_pg_sandbox.py` 这一类"验证隔离本身
|
||||
是否成立"的用例;业务断言一律走 `PgSandbox.dsn`。
|
||||
"""
|
||||
import asyncpg
|
||||
|
||||
dsn = _require_admin_dsn()
|
||||
|
||||
async def probe(sql: str, *args: object) -> list[tuple]:
|
||||
conn = await asyncpg.connect(dsn, timeout=10)
|
||||
try:
|
||||
return [tuple(r) for r in await conn.fetch(sql, *args)]
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
return probe
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def pg_sandbox():
|
||||
"""一次性沙箱工厂: `await pg_sandbox(ddl=..., role=...)`,清理由 fixture 兜底。
|
||||
|
||||
同一条用例可以要多个沙箱(如"A 的角色去动 B 的表"),它们按后进先出清理。
|
||||
"""
|
||||
import asyncpg
|
||||
|
||||
admin_dsn = _require_admin_dsn()
|
||||
# 清理动作栈: 每建成一个对象就入栈一条,setup 中途失败与正常 teardown 共用
|
||||
# 同一条退栈路径——两处各写一份的话,失败那条永远是没被测过的那份。
|
||||
cleanups: list[str] = []
|
||||
|
||||
async def _run_as_admin(*statements: str) -> None:
|
||||
conn = await asyncpg.connect(admin_dsn, timeout=10)
|
||||
try:
|
||||
for statement in statements:
|
||||
await conn.execute(statement)
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
async def _unwind(statements: list[str]) -> None:
|
||||
"""逆序执行清理并**逐条容错**: 一条失败不该拖累其余对象的清理。
|
||||
|
||||
吞掉异常是不行的(残留会静默累积),但让第一条失败中断整栈更糟——角色是
|
||||
全局对象,漏掉的每一个都要人手工去删。故全部试完再抛出第一个异常。
|
||||
"""
|
||||
first: BaseException | None = None
|
||||
for statement in reversed(statements):
|
||||
try:
|
||||
await _run_as_admin(statement)
|
||||
except Exception as exc: # noqa: BLE001 — 见 docstring: 收集而非吞没
|
||||
first = first or exc
|
||||
statements.clear()
|
||||
if first is not None:
|
||||
raise first
|
||||
|
||||
async def make(
|
||||
*,
|
||||
ddl: str | None = None,
|
||||
extra: Sequence[str] = (),
|
||||
role: _Role = "none",
|
||||
grants: Sequence[str] = ("SELECT", "INSERT"),
|
||||
) -> PgSandbox:
|
||||
# 权限门在建任何对象**之前**: pytest.skip 抛的是 BaseException,若它在
|
||||
# 已建对象之后触发,清理会去 DROP 从未建成的东西并把 skip 盖掉。
|
||||
if role != "none":
|
||||
conn = await asyncpg.connect(admin_dsn, timeout=10)
|
||||
try:
|
||||
can_create = await conn.fetchval(
|
||||
"SELECT rolcreaterole OR rolsuper FROM pg_roles WHERE rolname = current_user"
|
||||
)
|
||||
finally:
|
||||
await conn.close()
|
||||
if not can_create:
|
||||
pytest.skip("当前账号无权建临时角色,跳过需要独占角色的用例")
|
||||
|
||||
# schema 与角色的前缀有意不同: 同名会让 "$user" 命中自有 schema 并遮蔽
|
||||
# 共享表,于是"search_path 落到共享表"这个最坏情况就再也构造不出来。
|
||||
suffix = uuid4().hex[:12]
|
||||
schema = f"pgw_s_{suffix}"
|
||||
role_name = f"pgw_r_{suffix}" if role != "none" else None
|
||||
|
||||
# 本次调用自己的清理栈: 失败只回滚**本次**建成的对象。同一条用例常要两个
|
||||
# 沙箱(如"A 的角色去动 B 的表"),回滚整栈会把已通过断言依赖的对象也删掉。
|
||||
local: list[str] = []
|
||||
try:
|
||||
if role_name is not None:
|
||||
await _run_as_admin(f"CREATE ROLE {role_name} LOGIN PASSWORD '{_SANDBOX_PASSWORD}'")
|
||||
# DROP OWNED BY 必须排在 DROP ROLE 之前: 角色仍持有对象时删不掉
|
||||
local.append(f"DROP ROLE IF EXISTS {role_name}")
|
||||
local.append(f"DROP OWNED BY {role_name}")
|
||||
owner_clause = f" AUTHORIZATION {role_name}" if role == "owner" else ""
|
||||
await _run_as_admin(f"CREATE SCHEMA {schema}{owner_clause}")
|
||||
local.append(f"DROP SCHEMA IF EXISTS {schema} CASCADE")
|
||||
|
||||
bare = _as_role(admin_dsn, role_name) if role_name is not None else None
|
||||
# role="owner" 时 DDL 由角色自己执行,表属主才会是它;"grantee" 的现场
|
||||
# 恰恰相反——表由别的账号建好,角色只拿到表级权限。
|
||||
ddl_dsn = _with_search_path(bare if role == "owner" else admin_dsn, schema)
|
||||
if ddl is not None:
|
||||
conn = await asyncpg.connect(ddl_dsn, timeout=10)
|
||||
try:
|
||||
await conn.execute(ddl)
|
||||
for statement in extra:
|
||||
await conn.execute(statement)
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
if role == "grantee":
|
||||
await _run_as_admin(f"GRANT USAGE ON SCHEMA {schema} TO {role_name}")
|
||||
if ddl is not None:
|
||||
await _run_as_admin(
|
||||
f"GRANT {', '.join(grants)} ON ALL TABLES IN SCHEMA {schema} TO {role_name}"
|
||||
)
|
||||
# 关键: 绝不 GRANT CREATE ON SCHEMA —— 缺的正是这一项
|
||||
except BaseException:
|
||||
await _unwind(local)
|
||||
raise
|
||||
cleanups.extend(local)
|
||||
|
||||
used = bare if role_name is not None else admin_dsn
|
||||
return PgSandbox(
|
||||
schema=schema,
|
||||
role=role_name,
|
||||
dsn=_with_search_path(used, schema),
|
||||
bare_dsn=bare,
|
||||
)
|
||||
|
||||
yield make
|
||||
await _unwind(cleanups)
|
||||
@@ -0,0 +1,163 @@
|
||||
"""`conftest.py` 沙箱工厂自身的行为测试(issue #18 Task 1)。
|
||||
|
||||
工厂是本次一切隔离的地基: 它若在 setup 中途失败时漏掉清理、或让角色名与
|
||||
schema 名撞上,受害的不是这一个文件,而是此后每一条 PG 用例。故它必须先被测。
|
||||
|
||||
**这里的断言全部只看自建对象与 PG catalog**,不读任何共享数据。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
_DDL = "CREATE TABLE llm_calls (call_id TEXT PRIMARY KEY, created_at TIMESTAMPTZ DEFAULT now())"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def assert_no_leftovers(pg_catalog_probe):
|
||||
"""收集沙箱名,在 `pg_sandbox` 清理之后回查它们是否真的没了。
|
||||
|
||||
必须比 `pg_sandbox` **先** setup: pytest 的 finalizer 是后进先出,先 setup
|
||||
的后 teardown——本 fixture 的检查因此发生在沙箱清理之后,而不是之前。
|
||||
"""
|
||||
seen: list[tuple[str, str | None]] = []
|
||||
yield seen
|
||||
for schema, role in seen:
|
||||
left = await pg_catalog_probe("SELECT nspname FROM pg_namespace WHERE nspname = $1", schema)
|
||||
assert left == [], f"沙箱 schema 未清理: {schema}"
|
||||
if role is not None:
|
||||
left = await pg_catalog_probe("SELECT rolname FROM pg_roles WHERE rolname = $1", role)
|
||||
assert left == [], f"沙箱角色未清理: {role}"
|
||||
|
||||
|
||||
async def _oid_of_llm_calls(dsn: str) -> int | None:
|
||||
import asyncpg
|
||||
|
||||
conn = await asyncpg.connect(dsn, timeout=10)
|
||||
try:
|
||||
return await conn.fetchval("SELECT to_regclass('llm_calls')::oid")
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
|
||||
class TestSchemaOnlySandbox:
|
||||
async def test_table_lands_in_the_sandbox_schema_and_bare_dsn_is_absent(self, pg_sandbox):
|
||||
"""`role="none"`: 表落在自建 schema 下;不发角色,故没有裸 DSN 可给。"""
|
||||
sandbox = await pg_sandbox(ddl=_DDL)
|
||||
|
||||
import asyncpg
|
||||
|
||||
conn = await asyncpg.connect(sandbox.dsn, timeout=10)
|
||||
try:
|
||||
where = await conn.fetchval(
|
||||
"SELECT n.nspname FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace "
|
||||
"WHERE c.oid = to_regclass('llm_calls')"
|
||||
)
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
assert where == sandbox.schema
|
||||
assert sandbox.role is None
|
||||
assert sandbox.bare_dsn is None
|
||||
|
||||
|
||||
class TestOwnerRoleSandbox:
|
||||
async def test_the_role_owns_its_own_table(self, pg_sandbox):
|
||||
"""`role="owner"`: 表由角色自己建,故属主是它——与"用维护角色跑"的现场一致。"""
|
||||
sandbox = await pg_sandbox(ddl=_DDL, role="owner")
|
||||
|
||||
import asyncpg
|
||||
|
||||
conn = await asyncpg.connect(sandbox.dsn, timeout=10)
|
||||
try:
|
||||
owner = await conn.fetchval(
|
||||
"SELECT pg_get_userbyid(relowner) FROM pg_class WHERE oid = to_regclass('llm_calls')"
|
||||
)
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
assert owner == sandbox.role
|
||||
# 名字必须错开: 同名会让 "$user" 命中自有 schema 并遮蔽真表,
|
||||
# 最坏情况用例就再也走不到那条真实路径上(设计 §5.1 实测)
|
||||
assert sandbox.role != sandbox.schema
|
||||
assert not sandbox.role.startswith("pgw_s_")
|
||||
assert not sandbox.schema.startswith("pgw_r_")
|
||||
|
||||
async def test_bare_dsn_falls_through_to_the_default_search_path(self, pg_sandbox):
|
||||
"""裸 DSN 必须真的回落到 `"$user", public`——最坏情况用例全靠它构造现场。"""
|
||||
sandbox = await pg_sandbox(ddl=_DDL, role="owner")
|
||||
|
||||
import asyncpg
|
||||
|
||||
conn = await asyncpg.connect(sandbox.bare_dsn, timeout=10)
|
||||
try:
|
||||
path = await conn.fetchval("SHOW search_path")
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
assert path == '"$user", public'
|
||||
# 裸 DSN 解析到的绝不能是沙箱里那张表(否则"落到共享表"的现场是假的)
|
||||
assert await _oid_of_llm_calls(sandbox.bare_dsn) != await _oid_of_llm_calls(sandbox.dsn)
|
||||
|
||||
|
||||
class TestGranteeRoleSandbox:
|
||||
async def test_grantee_can_write_but_cannot_create(self, pg_sandbox):
|
||||
"""`role="grantee"`: 表属主是 admin,角色只拿表级权限——最小权限部署的现场。"""
|
||||
import asyncpg
|
||||
|
||||
sandbox = await pg_sandbox(ddl=_DDL, role="grantee")
|
||||
|
||||
conn = await asyncpg.connect(sandbox.dsn, timeout=10)
|
||||
try:
|
||||
await conn.execute("INSERT INTO llm_calls (call_id) VALUES ('g1')")
|
||||
assert await conn.fetchval("SELECT count(*) FROM llm_calls") == 1
|
||||
with pytest.raises(asyncpg.exceptions.InsufficientPrivilegeError):
|
||||
await conn.execute("CREATE TABLE another (x TEXT)")
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
|
||||
class TestCleanup:
|
||||
async def test_setup_failure_leaves_nothing_behind(self, pg_sandbox, pg_catalog_probe):
|
||||
"""建到一半失败时也必须删净——角色是**全局**对象,残留不随库消失。"""
|
||||
before_schemas = await pg_catalog_probe(
|
||||
"SELECT nspname FROM pg_namespace WHERE nspname LIKE 'pgw!_%' ESCAPE '!'"
|
||||
)
|
||||
before_roles = await pg_catalog_probe(
|
||||
"SELECT rolname FROM pg_roles WHERE rolname LIKE 'pgw!_%' ESCAPE '!'"
|
||||
)
|
||||
|
||||
with pytest.raises(Exception): # noqa: B017 — 工厂原样抛出 PG 的 DDL 错误
|
||||
await pg_sandbox(ddl="CREATE TABLE llm_calls (bad NOT_A_REAL_TYPE)", role="owner")
|
||||
|
||||
assert (
|
||||
await pg_catalog_probe(
|
||||
"SELECT nspname FROM pg_namespace WHERE nspname LIKE 'pgw!_%' ESCAPE '!'"
|
||||
)
|
||||
== before_schemas
|
||||
)
|
||||
assert (
|
||||
await pg_catalog_probe(
|
||||
"SELECT rolname FROM pg_roles WHERE rolname LIKE 'pgw!_%' ESCAPE '!'"
|
||||
)
|
||||
== before_roles
|
||||
)
|
||||
|
||||
async def test_a_failure_does_not_roll_back_earlier_sandboxes(self, pg_sandbox):
|
||||
"""一次失败只回滚它自己建的东西——同一条用例里先建成的沙箱必须毫发无损。
|
||||
|
||||
"A 的角色去动 B 的表"这类用例一条要两个沙箱;若失败回滚把整栈清空,受害的
|
||||
是那些**已经通过**的断言所依赖的对象,而症状会以"表不见了"的形态出现在
|
||||
与真因无关的地方。
|
||||
"""
|
||||
good = await pg_sandbox(ddl=_DDL, role="owner")
|
||||
|
||||
with pytest.raises(Exception): # noqa: B017 — 工厂原样抛出 PG 的 DDL 错误
|
||||
await pg_sandbox(ddl="CREATE TABLE llm_calls (bad NOT_A_REAL_TYPE)", role="owner")
|
||||
|
||||
assert await _oid_of_llm_calls(good.dsn) is not None, "先前建成的沙箱被误清理"
|
||||
|
||||
async def test_teardown_removes_schema_and_role(self, assert_no_leftovers, pg_sandbox):
|
||||
"""正常路径的清理: 断言发生在 `pg_sandbox` teardown **之后**(见 fixture 说明)。"""
|
||||
sandbox = await pg_sandbox(ddl=_DDL, role="owner")
|
||||
assert_no_leftovers.append((sandbox.schema, sandbox.role))
|
||||
Reference in New Issue
Block a user