honcho/tests/test_dependencies.py

import uuid
from typing import Any

import pytest

import src.dependencies as dependencies_module
from src.config import settings
from src.db import request_context
from src.dependencies import get_db as real_get_db
from src.dependencies import tracked_db as real_tracked_db


class FakeSession:
    def __init__(self, *, in_transaction: bool = False):
        self._in_transaction: bool = in_transaction
        self.execute_calls: list[tuple[Any, ...]] = []
        self.rollback_calls: int = 0
        self.close_calls: int = 0
        self.connection_calls: int = 0

    async def connection(self) -> None:
        # Tracks checkout attempts so tests can assert get_db/tracked_db stay
        # lazy (they should never force a checkout themselves).
        self.connection_calls += 1

    async def execute(self, statement: Any, params: Any = None) -> None:
        self.execute_calls.append((statement, params))

    async def rollback(self) -> None:
        self.rollback_calls += 1

    async def close(self) -> None:
        self.close_calls += 1

    def in_transaction(self) -> bool:
        return self._in_transaction


@pytest.mark.asyncio
async def test_get_db_yields_lazily_without_checkout_or_tracing(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    # get_db must NOT touch the connection or run set_config itself — checkout
    # happens lazily inside the AsyncSession on first DB use, so a handler doing
    # non-DB work before its first query never pins a connection.
    fake_db = FakeSession()
    monkeypatch.setattr(dependencies_module, "SessionLocal", lambda: fake_db)
    monkeypatch.setattr(settings.DB, "TRACING", True)  # still no set_config here

    dep_gen = real_get_db()
    try:
        db = await anext(dep_gen)
        assert db is fake_db
        assert fake_db.connection_calls == 0  # no eager checkout
        assert fake_db.execute_calls == []  # no set_config in get_db
    finally:
        await dep_gen.aclose()

    assert fake_db.rollback_calls == 1  # unconditional rollback in finally
    assert fake_db.close_calls == 1


@pytest.mark.asyncio
async def test_get_db_rolls_back_and_closes_when_consumer_raises(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    fake_db = FakeSession()
    monkeypatch.setattr(dependencies_module, "SessionLocal", lambda: fake_db)
    monkeypatch.setattr(settings.DB, "TRACING", False)

    dep_gen = real_get_db()
    await anext(dep_gen)

    with pytest.raises(RuntimeError, match="boom"):
        await dep_gen.athrow(RuntimeError("boom"))

    assert fake_db.rollback_calls == 2  # once in except, once in finally
    assert fake_db.close_calls == 1


@pytest.mark.asyncio
async def test_tracked_db_creates_and_resets_task_context(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    fake_db = FakeSession()
    monkeypatch.setattr(dependencies_module, "SessionLocal", lambda: fake_db)
    monkeypatch.setattr(
        uuid,
        "uuid4",
        lambda: uuid.UUID("12345678-1234-5678-1234-567812345678"),
    )

    clear_token = request_context.set(None)
    try:
        async with real_tracked_db("cleanup_job"):
            # tracked_db sets the task context so the lazy session can read it.
            assert request_context.get() == "task:cleanup_job:12345678"
    finally:
        request_context.reset(clear_token)

    assert request_context.get() is None
    assert fake_db.rollback_calls == 1  # unconditional rollback in finally
    assert fake_db.close_calls == 1


@pytest.mark.asyncio
async def test_tracked_db_preserves_existing_request_context(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    fake_db = FakeSession()
    monkeypatch.setattr(dependencies_module, "SessionLocal", lambda: fake_db)

    context_token = request_context.set("request:existing")
    try:
        async with real_tracked_db("ignored_op"):
            assert request_context.get() == "request:existing"
    finally:
        request_context.reset(context_token)

    assert fake_db.rollback_calls == 1  # unconditional rollback in finally
    assert fake_db.close_calls == 1


@pytest.mark.asyncio
async def test_tracked_db_rolls_back_on_error_and_closes(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    fake_db = FakeSession()
    monkeypatch.setattr(dependencies_module, "SessionLocal", lambda: fake_db)
    monkeypatch.setattr(settings.DB, "TRACING", False)

    with pytest.raises(ValueError, match="failed operation"):
        async with real_tracked_db("operation"):
            raise ValueError("failed operation")

    assert fake_db.rollback_calls == 2  # once in except, once in finally
    assert fake_db.close_calls == 1


@pytest.mark.asyncio
async def test_tracked_db_rolls_back_open_transaction_on_exit(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    fake_db = FakeSession(in_transaction=True)
    monkeypatch.setattr(dependencies_module, "SessionLocal", lambda: fake_db)
    monkeypatch.setattr(settings.DB, "TRACING", False)

    async with real_tracked_db("operation"):
        pass

    assert fake_db.rollback_calls == 1
    assert fake_db.close_calls == 1


@pytest.mark.asyncio
async def test_tracked_db_read_only_uses_read_sessionmaker(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    # read_only=True must construct the session from ReadSessionLocal (the
    # AUTOCOMMIT engine) — and never touch SessionLocal — while keeping the
    # same rollback/close teardown.
    read_fake = FakeSession()
    monkeypatch.setattr(dependencies_module, "ReadSessionLocal", lambda: read_fake)
    monkeypatch.setattr(
        dependencies_module,
        "SessionLocal",
        lambda: pytest.fail("read_only window constructed a write session"),
    )

    async with real_tracked_db("read_op", read_only=True) as db:
        assert db is read_fake

    assert read_fake.rollback_calls == 1
    assert read_fake.close_calls == 1


@pytest.mark.asyncio
async def test_get_read_db_rolls_back_and_closes(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    read_fake = FakeSession()
    monkeypatch.setattr(dependencies_module, "ReadSessionLocal", lambda: read_fake)

    dep_gen = dependencies_module.get_read_db()
    try:
        db = await anext(dep_gen)
        assert db is read_fake
        assert read_fake.connection_calls == 0  # still lazy, no eager checkout
    finally:
        await dep_gen.aclose()

    assert read_fake.rollback_calls == 1
    assert read_fake.close_calls == 1


def test_read_engine_is_autocommit_and_shares_pool() -> None:
    # The read engine must differ from the write engine ONLY by isolation
    # level: AUTOCOMMIT (so reads never autobegin a transaction) on the same
    # underlying pool (no second connection budget).
    from src.db import engine, read_engine

    assert (
        read_engine.sync_engine._execution_options.get(  # pyright: ignore[reportPrivateUsage]
            "isolation_level"
        )
        == "AUTOCOMMIT"
    )
    assert read_engine.sync_engine.pool is engine.sync_engine.pool


@pytest.mark.asyncio
async def test_read_only_session_runs_in_autocommit_on_the_wire() -> None:
    # Wire-level guarantee behind the whole read-path fix: a read_only session's
    # connection has the DBAPI autocommit flag set, so psycopg emits no BEGIN
    # and the backend sits in state 'idle' (not 'idle in transaction') after a
    # statement returns. NOTE: get_isolation_level() can NOT verify this — it
    # reports the server's transaction_isolation GUC (READ COMMITTED), because
    # autocommit is a driver behavior, not a server isolation level.
    from sqlalchemy import text

    from src.db import read_engine

    async with real_tracked_db("read_op", read_only=True) as db:
        pid = (await db.execute(text("SELECT pg_backend_pid()"))).scalar()
        conn = await db.connection()
        raw = (await conn.get_raw_connection()).driver_connection
        assert raw is not None
        assert raw.autocommit is True

        # Definitive check, from a second connection: after the SELECT above,
        # the session's backend must be plain 'idle' — an open transaction
        # would report 'idle in transaction' and be reapable in production.
        async with read_engine.connect() as observer:
            state = (
                await observer.execute(
                    text("SELECT state FROM pg_stat_activity WHERE pid = :p"),
                    {"p": pid},
                )
            ).scalar()
        assert state == "idle"


@pytest.mark.asyncio
async def test_write_session_holds_idle_in_transaction_after_select() -> None:
    # Contrast guard documenting WHY the read engine exists: the default
    # (transactional) session autobegins on the first statement and leaves the
    # backend 'idle in transaction' until rollback/close — the state that
    # Postgres's idle_in_transaction_session_timeout reaps and that pins a
    # transaction-mode pooler backend.
    from sqlalchemy import text

    from src.db import read_engine

    async with real_tracked_db("write_op") as db:
        pid = (await db.execute(text("SELECT pg_backend_pid()"))).scalar()
        async with read_engine.connect() as observer:
            state = (
                await observer.execute(
                    text("SELECT state FROM pg_stat_activity WHERE pid = :p"),
                    {"p": pid},
                )
            ).scalar()
        assert state == "idle in transaction"


@pytest.mark.asyncio
async def test_read_only_session_works_with_tracing_checkout_hook() -> None:
    # Regression: the DB.TRACING checkout hook runs set_config() at pool
    # checkout, BEFORE the dialect applies the read engine's AUTOCOMMIT
    # isolation level. If that statement is allowed to autobegin a transaction,
    # psycopg then refuses to switch the connection into AUTOCOMMIT
    # ("can't change 'autocommit' now: connection in transaction") and every
    # read_only session 500s under TRACING. The hook must run in autocommit so
    # it leaves the connection idle. This combination is otherwise untested
    # because DB.TRACING defaults to false.
    from sqlalchemy import event, text

    from src.db import (
        _set_application_name_on_checkout,  # pyright: ignore[reportPrivateUsage]
        engine,
        read_engine,
    )

    context_token = request_context.set("tracing-regression")
    event.listen(engine.sync_engine, "checkout", _set_application_name_on_checkout)
    try:
        async with real_tracked_db("read_op", read_only=True) as db:
            pid = (await db.execute(text("SELECT pg_backend_pid()"))).scalar()
            app_name = (await db.execute(text("SHOW application_name"))).scalar()
            conn = await db.connection()
            raw = (await conn.get_raw_connection()).driver_connection
            assert raw is not None
            # AUTOCOMMIT was applied despite the checkout hook running first.
            assert raw.autocommit is True
            # The hook still tagged the connection (set_config is session-scoped,
            # so it survives the autocommit boundary).
            assert app_name == "tracing-regression"
            # Backend is idle, not idle-in-transaction: the no-BEGIN guarantee
            # holds even with the hook firing.
            async with read_engine.connect() as observer:
                state = (
                    await observer.execute(
                        text("SELECT state FROM pg_stat_activity WHERE pid = :p"),
                        {"p": pid},
                    )
                ).scalar()
            assert state == "idle"
    finally:
        event.remove(engine.sync_engine, "checkout", _set_application_name_on_checkout)
        request_context.reset(context_token)