import psycopg
import pytest

from lyra_worker.db import connect, wait_for_db


def test_terminated_backend_raises_operational_error_then_reconnects(conn):
    """The run_forever reconnect boundary catches psycopg.OperationalError and re-runs
    wait_for_db(). This exercises both primitives: a server-side backend termination (the
    real AdminShutdown case) surfaces as OperationalError, and wait_for_db reconnects."""
    victim = connect()
    try:
        with victim.cursor() as cur:
            cur.execute("SELECT pg_backend_pid()")
            pid = cur.fetchone()[0]
        with conn.cursor() as cur:  # terminate the victim's backend, as a DB restart would
            cur.execute("SELECT pg_terminate_backend(%s)", (pid,))
        conn.commit()
        with pytest.raises(psycopg.OperationalError):  # the exact type the loop catches
            with victim.cursor() as cur:
                cur.execute("SELECT 1")
                cur.fetchone()
    finally:
        try:
            victim.close()
        except Exception:
            pass

    fresh = wait_for_db()  # reconnect exactly as run_forever does
    try:
        with fresh.cursor() as cur:
            cur.execute("SELECT 1")
            assert cur.fetchone()[0] == 1
    finally:
        fresh.close()