Lamont/tests/integration/conftest.py

"""Isolated PostgreSQL and Flask fixtures for integration tests."""

import os
from pathlib import Path
import re
import uuid

import psycopg2
from psycopg2 import sql
import pytest
from dotenv import load_dotenv


PROJECT_ROOT = Path(__file__).resolve().parents[2]
load_dotenv(PROJECT_ROOT / ".env", override=False)

TEST_DATABASE_PREFIX = "llm_bot_integration_"
SERVICE_KEY = "integration-service-key-0123456789abcdef"
ALLOWED_DISCORD_ID = "integration-discord-user"


def _databaseConfig(databaseName=None):
    return {
        "host": os.environ.get("DB_HOST", "localhost"),
        "port": int(os.environ.get("DB_PORT", 5432)),
        "dbname": databaseName or os.environ.get("DB_NAME", "app"),
        "user": os.environ.get("DB_USER", "app"),
        "password": os.environ.get("DB_PASS", ""),
        "connect_timeout": 3,
    }


def _databaseUnavailable(message):
    if os.environ.get("CI", "").lower() == "true":
        pytest.fail(message)
    pytest.skip(message)


@pytest.fixture(scope="session")
def postgresServer():
    """Skip once locally, but fail CI when its required PostgreSQL is absent."""
    sourceDatabase = os.environ.get("DB_NAME", "app")
    try:
        connection = psycopg2.connect(**_databaseConfig(sourceDatabase))
    except psycopg2.OperationalError as error:
        _databaseUnavailable(
            f"PostgreSQL integration database is unavailable: {error}"
        )

    with connection:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                SELECT rolsuper OR rolcreatedb AS can_create_database
                FROM pg_roles
                WHERE rolname = CURRENT_USER
                """
            )
            canCreateDatabase = cursor.fetchone()[0]
    connection.close()
    if not canCreateDatabase:
        _databaseUnavailable(
            "PostgreSQL integration user requires CREATEDB for isolated tests"
        )
    return sourceDatabase


@pytest.fixture
def isolatedDatabase(monkeypatch, postgresServer):
    """Create a disposable database without altering the configured database."""
    databaseName = f"{TEST_DATABASE_PREFIX}{uuid.uuid4().hex}"
    assert re.fullmatch(r"llm_bot_integration_[0-9a-f]{32}", databaseName)

    adminConnection = psycopg2.connect(**_databaseConfig(postgresServer))

    adminConnection.autocommit = True
    try:
        with adminConnection.cursor() as cursor:
            cursor.execute(
                sql.SQL("CREATE DATABASE {}").format(sql.Identifier(databaseName))
            )
    except psycopg2.Error as error:
        adminConnection.close()
        _databaseUnavailable(
            f"PostgreSQL user cannot create an isolated database: {error}"
        )

    monkeypatch.setenv("DB_NAME", databaseName)
    try:
        yield databaseName
    finally:
        # Every production helper closes its connection, but terminate any failed-test
        # leftovers before dropping only the uniquely named test database.
        with adminConnection.cursor() as cursor:
            cursor.execute(
                """
                SELECT pg_terminate_backend(pid)
                FROM pg_stat_activity
                WHERE datname = %s AND pid != pg_backend_pid()
                """,
                (databaseName,),
            )
            cursor.execute(
                sql.SQL("DROP DATABASE {}").format(sql.Identifier(databaseName))
            )
        adminConnection.close()


@pytest.fixture
def migratedDatabase(isolatedDatabase):
    from core.migrations import upgrade

    applied = upgrade()
    assert applied
    return isolatedDatabase


@pytest.fixture
def app(migratedDatabase, monkeypatch):
    monkeypatch.setenv("JWT_SECRET", "integration-only-jwt-secret")
    monkeypatch.setenv("BOT_API_KEY", SERVICE_KEY)
    monkeypatch.setenv(
        "BOT_API_KEY_SCOPES",
        "discord:session,outbox:claim,outbox:deliver",
    )
    monkeypatch.setenv("DISCORD_ENROLLMENT_MODE", "allowlist")
    monkeypatch.setenv("DISCORD_ALLOWLIST", ALLOWED_DISCORD_ID)
    monkeypatch.setenv("DEFAULT_TIMEZONE", "UTC")

    from api.main import createApp

    flaskApp = createApp()
    flaskApp.config.update(TESTING=True)
    return flaskApp


@pytest.fixture
def client(app):
    return app.test_client()


@pytest.fixture
def createAuthenticatedUser(client):
    created = 0

    def create(username=None, password="correct-horse-battery-staple"):
        nonlocal created
        created += 1
        username = username or f"integration-user-{created}-{uuid.uuid4().hex}"
        registration = client.post(
            "/api/register",
            json={"username": username, "password": password, "timezone": "UTC"},
        )
        assert registration.status_code == 201, registration.get_json()

        login = client.post(
            "/api/login", json={"username": username, "password": password}
        )
        assert login.status_code == 200, login.get_json()
        token = login.get_json()["token"]

        uuidResponse = client.get(
            f"/api/getUserUUID/{username}",
            headers={"Authorization": f"Bearer {token}"},
        )
        assert uuidResponse.status_code == 200, uuidResponse.get_json()
        return {
            "username": username,
            "password": password,
            "token": token,
            "user_uuid": uuidResponse.get_json(),
        }

    return create