PolyGateway/tests/integration/test_pg_sandbox.py

"""`conftest.py` 沙箱工厂自身的行为测试(issue #18 Task 1)。

工厂是本次一切隔离的地基: 它若在 setup 中途失败时漏掉清理、或让角色名与
schema 名撞上,受害的不是这一个文件,而是此后每一条 PG 用例。故它必须先被测。

**这里的断言全部只看自建对象与 PG catalog**,不读任何共享数据。
"""

from __future__ import annotations

import pytest

from tests.integration.conftest import _as_role

_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 TestRoleDsnConstruction:
    """凭据替换失败必须**当场报错**,不许退回管理身份(合并前审查的 P1)。

    这条防线的失效形态特别隐蔽: 替换不上时 `re.sub` 原样返回管理连接串,沙箱
    "看起来"建好了、用例照常绿,而 `bare_dsn` 其实是超级用户——最坏情况用例
    会拿它跑真实 `--apply`,把共享表删空之后才在 `assert returncode == 2` 上红。
    行已经没了。设计 §5.1 要的是"越界做不到",不是"越界会被发现"。
    """

    def test_inline_credentials_are_replaced(self):
        swapped = _as_role("postgresql://app:secret@h:5432/polygateway", "pgw_r_x")

        assert swapped.startswith("postgresql://pgw_r_x:")
        assert "app:secret" not in swapped

    @pytest.mark.parametrize(
        "dsn",
        [
            "postgresql://h:5432/polygateway",  # 口令走 PGPASSWORD / .pgpass / trust
            "postgresql:///polygateway?host=/var/run/postgresql",  # unix socket
        ],
    )
    def test_a_dsn_without_inline_credentials_is_refused(self, dsn):
        """这两种都是合法 DSN,今天的 .env 恰好不是它们——恰好而已。"""
        with pytest.raises(RuntimeError, match="沙箱角色"):
            _as_role(dsn, "pgw_r_x")


class TestCatalogProbeIsReadOnly:
    """探针拿的是管理连接,故它只许查——否则"用例够不到管理能力"就是句空话。"""

    async def test_non_select_statements_are_refused(self, pg_catalog_probe):
        with pytest.raises(RuntimeError, match="只接受 SELECT"):
            await pg_catalog_probe("DELETE FROM llm_calls WHERE call_id = 'nope'")


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))