magnus919_agent-skills/postgres/tests/test_pgdiag.py

#!/usr/bin/env python3
"""Deterministic tests for the postgres/scripts/pgdiag diagnostic tool.

Runs the script as a subprocess so the tests exercise the real CLI surface
(--help, --json, --check selection, --plan-for, exit codes, JSON payloads).
No PostgreSQL server is required: a fake psql stub is written to a temp
directory at test time and pointed at with --psql. The stub also enforces
the read-only contract by failing any invocation that does not carry the
default_transaction_read_only session setting.
"""
import json
import os
import stat
import subprocess
import sys
import tempfile
import unittest
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
SCRIPT = ROOT / "scripts" / "pgdiag"

STUB_TEMPLATE = """\
#!/usr/bin/env python3
import os
import sys

args = sys.argv[1:]
sqls = []
index = 0
while index < len(args):
    if args[index] == "-c":
        sqls.append(args[index + 1])
        index += 2
    else:
        index += 1

if not any("default_transaction_read_only" in item for item in sqls):
    print("read-only session setting missing", file=sys.stderr)
    sys.exit(1)

sql = sqls[-1]

if os.environ.get("PGDIAG_STUB_FAIL") == "1":
    print("injected psql failure", file=sys.stderr)
    sys.exit(1)

if "server_version" in sql:
    print("mydb|dba|16.4|160004|f|2026-07-01 08:00:00+00")
elif "pg_stat_archiver" in sql:
    print("42|0|00000001000000000000002A|2026-07-01 00:00:00+00||")
elif "application_name" in sql:
    print("standby-1|streaming|sync|10.0.0.5|0/2A000000|0/2A000000|0/2A000000|0/2A000000")
elif "pg_last_wal_receive_lsn" in sql:
    print("f|0/2A000000|0/2A000000")
elif "pg_extension" in sql:
    print("plpgsql|1.0")
elif "idx_scan DESC" in sql:
    print("public|orders|orders_pkey|12345|24680|12000")
elif "idx_scan = 0" in sql:
    print("public|events|events_created_at_idx|0")
elif "pg_index" in sql:
    pass
elif "seq_tup_read DESC" in sql:
    print("public|orders|5000|1500000|120")
elif "n_dead_tup > 0" in sql:
    print("public|orders|100000|5000|4.8")
elif "pg_stat_activity" in sql:
    print("active|12")
    print("idle|38")
elif "current_setting" in sql:
    print("max_connections|100")
    print("wal_level|replica")
    print("archive_mode|on")
elif "pg_database_size" in sql:
    print("mydb|128 MB")
elif "EXPLAIN" in sql:
    print('[{"Plan": {"Node Type": "Seq Scan", "Plan Rows": 1000}}]')
elif "current_database()" in sql:
    print("mydb|dba|16.4|160004|f|2026-07-01 08:00:00+00")
else:
    pass
"""


def write_stub(directory: Path) -> Path:
    stub = directory / "fake-psql"
    stub.write_text(STUB_TEMPLATE, encoding="utf-8")
    stub.chmod(stub.stat().st_mode | stat.S_IXUSR)
    return stub


def run_script(*args: str, fail: bool = False) -> subprocess.CompletedProcess:
    env = os.environ.copy()
    if fail:
        env["PGDIAG_STUB_FAIL"] = "1"
    with tempfile.TemporaryDirectory() as tmp:
        stub = write_stub(Path(tmp))
        proc = subprocess.run(
            [sys.executable, str(SCRIPT), "--psql", str(stub), *args],
            capture_output=True,
            text=True,
            env=env,
            timeout=30,
        )
    return proc


class HelpTests(unittest.TestCase):
    def test_help_exits_zero_without_cluster(self):
        proc = subprocess.run(
            [sys.executable, str(SCRIPT), "--help"],
            capture_output=True,
            text=True,
            timeout=30,
        )
        self.assertEqual(proc.returncode, 0)
        self.assertIn("--json", proc.stdout)
        self.assertIn("read-only", proc.stdout.lower())
        self.assertIn("default_transaction_read_only", proc.stdout)

    def test_version_flag(self):
        proc = subprocess.run(
            [sys.executable, str(SCRIPT), "--version"],
            capture_output=True,
            text=True,
            timeout=30,
        )
        self.assertEqual(proc.returncode, 0)
        self.assertIn("pgdiag", proc.stdout)


class JsonRunTests(unittest.TestCase):
    def test_full_json_run_is_parseable(self):
        proc = run_script("--json")
        self.assertEqual(proc.returncode, 0, proc.stderr)
        payload = json.loads(proc.stdout)
        self.assertTrue(payload["ok"])
        self.assertEqual(payload["server"]["version"], "16.4")
        self.assertEqual(payload["server"]["database"], "mydb")
        names = [check["name"] for check in payload["checks"]]
        self.assertIn("identity", names)
        self.assertIn("config", names)
        self.assertIn("index_usage", names)
        self.assertIn("bloat", names)
        self.assertIn("wal_archive", names)
        self.assertIn("replication", names)
        self.assertIn("extensions", names)
        self.assertIn("databases", names)
        for check in payload["checks"]:
            self.assertEqual(check["status"], "ok", check["name"])

    def test_text_output_mentions_server_and_checks(self):
        proc = run_script()
        self.assertEqual(proc.returncode, 0, proc.stderr)
        self.assertIn("Server identity and version", proc.stdout)
        self.assertIn("WAL archiving health", proc.stdout)

    def test_check_subset_selection(self):
        proc = run_script("--check", "identity", "--check", "config", "--json")
        self.assertEqual(proc.returncode, 0, proc.stderr)
        payload = json.loads(proc.stdout)
        self.assertEqual([check["name"] for check in payload["checks"]], ["identity", "config"])

    def test_unknown_check_is_usage_error(self):
        proc = run_script("--check", "bogus", "--json")
        self.assertEqual(proc.returncode, 2)
        self.assertIn("unknown check", proc.stderr)

    def test_psql_error_recorded_per_check_but_run_continues(self):
        proc = run_script("--json", fail=True)
        self.assertEqual(proc.returncode, 1)
        payload = json.loads(proc.stdout)
        self.assertFalse(payload["ok"])
        names = [check["name"] for check in payload["checks"]]
        self.assertEqual(len(names), len(set(names)))
        self.assertTrue(all(check["status"] == "error" for check in payload["checks"]))


class BinaryAvailabilityTests(unittest.TestCase):
    def test_missing_psql_returns_127_with_json_error(self):
        proc = subprocess.run(
            [sys.executable, str(SCRIPT), "--psql", "/nonexistent/psql", "--json"],
            capture_output=True,
            text=True,
            timeout=30,
        )
        self.assertEqual(proc.returncode, 127)
        payload = json.loads(proc.stdout)
        self.assertFalse(payload["ok"])
        self.assertIn("psql binary not found", payload["error"])


class PlanForTests(unittest.TestCase):
    def test_plan_for_single_select_included(self):
        proc = run_script("--plan-for", "SELECT 1", "--json")
        self.assertEqual(proc.returncode, 0, proc.stderr)
        payload = json.loads(proc.stdout)
        plan_check = payload["checks"][-1]
        self.assertEqual(plan_check["name"], "plan")
        self.assertEqual(plan_check["status"], "ok")
        self.assertIsInstance(plan_check["plan"], list)

    def test_plan_for_rejects_write_statement(self):
        proc = run_script("--plan-for", "INSERT INTO t VALUES (1)", "--json")
        self.assertEqual(proc.returncode, 2)
        self.assertIn("read-only statement", proc.stdout)

    def test_plan_for_rejects_multiple_statements(self):
        proc = run_script("--plan-for", "SELECT 1; SELECT 2", "--json")
        self.assertEqual(proc.returncode, 2)


if __name__ == "__main__":
    unittest.main()