content-find-agent-new/tests/test_pattern_pg.py

from __future__ import annotations

import sys
import types

import pytest

from content_agent.errors import ContentAgentError, ErrorCode
from content_agent.integrations.pattern_pg import PatternPgClient


class FakeCursor:
    def __init__(self, *, fail_on_execute: bool = False) -> None:
        self.fail_on_execute = fail_on_execute
        self.executed: list[tuple[str, object | None]] = []

    def execute(self, sql, params=None):
        self.executed.append((sql, params))
        if self.fail_on_execute:
            raise RuntimeError('syntax error at or near "$1"')

    def fetchone(self):
        return (1,)


class FakeConnection:
    def __init__(self, cursor: FakeCursor) -> None:
        self._cursor = cursor
        self.closed = False

    def cursor(self):
        return self._cursor

    def close(self):
        self.closed = True


def _install_fake_pg8000(monkeypatch, connect_fn):
    fake_pkg = types.ModuleType("pg8000")
    fake_pkg.__path__ = []
    fake_dbapi = types.ModuleType("pg8000.dbapi")
    fake_dbapi.connect = connect_fn
    fake_pkg.dbapi = fake_dbapi
    monkeypatch.setitem(sys.modules, "pg8000", fake_pkg)
    monkeypatch.setitem(sys.modules, "pg8000.dbapi", fake_dbapi)


def _client() -> PatternPgClient:
    return PatternPgClient(
        host="127.0.0.1",
        port=5432,
        user="u",
        password="p",
        database="open_aigc",
        timeout_seconds=30,
    )


def test_pattern_pg_sets_statement_timeout_without_parameter_placeholder(monkeypatch):
    cursor = FakeCursor()
    connection = FakeConnection(cursor)
    _install_fake_pg8000(monkeypatch, lambda **kwargs: connection)

    assert _client().has_terminal_element(581, [101, 102]) is True

    timeout_sql, timeout_params = cursor.executed[0]
    assert timeout_sql == "SET statement_timeout = 30000"
    assert timeout_params is None
    assert "$1" not in timeout_sql
    assert connection.closed is True


def test_pattern_pg_query_failure_keeps_original_error_message(monkeypatch):
    cursor = FakeCursor(fail_on_execute=True)
    connection = FakeConnection(cursor)
    _install_fake_pg8000(monkeypatch, lambda **kwargs: connection)

    with pytest.raises(ContentAgentError) as exc:
        _client().has_terminal_element(581, [101])

    assert exc.value.error_code == ErrorCode.DB_SCHEMA_NOT_READY
    assert exc.value.detail["exception_type"] == "RuntimeError"
    assert 'syntax error at or near "$1"' in exc.value.detail["error_message"]
    assert exc.value.detail["execution_id"] == 581
    assert exc.value.detail["category_id_count"] == 1
    assert connection.closed is True