agentops/app/api/tests/_conftest/supabase.py

import os
import subprocess
from pathlib import Path
import pytest
from .common import is_github_actions, get_free_port


# these vars are also hard-coded inside the `supabase_server` Dockerfile and the
# github actions workflow
TEST_SUPABASE_HOST = "localhost"
TEST_SUPABASE_PORT = get_free_port(5432)
TEST_SUPABASE_NAME = "test_db"
TEST_SUPABASE_USER = "postgres"
TEST_SUPABASE_PASSWORD = "postgres"

SUPABASE_DOCKER_IMAGE_PATH = Path(__file__).parent / "supabase_server" / "Dockerfile"
SUPABASE_DOCKER_CONTAINER_NAME = "agentops-test-supabase"

__all__ = [
    'supabase_verify_test_environment',
    'supabase_setup_db_server',
    'db_session',
    'orm_session',
]


@pytest.fixture(scope="session", autouse=True)
def supabase_verify_test_environment():
    """Verify we don't have access to production credentials before running tests."""
    message = "%s environment variable is set! This risks connecting to production."
    assert not os.environ.get('SUPABASE_HOST'), message % "SUPABASE_HOST"
    assert not os.environ.get('SUPABASE_PORT'), message % "SUPABASE_PORT"
    assert not os.environ.get('SUPABASE_DATABASE'), message % "SUPABASE_DATABASE"
    assert not os.environ.get('SUPABASE_USER'), message % "SUPABASE_USER"
    assert not os.environ.get('SUPABASE_PASSWORD'), message % "SUPABASE_PASSWORD"


def supabase_start_docker():
    """Start the Docker container for PostgreSQL."""
    print("Starting Supabase/Postgres Docker container...")

    subprocess.run(['docker', 'stop', SUPABASE_DOCKER_CONTAINER_NAME], check=False)
    subprocess.run(['docker', 'rm', SUPABASE_DOCKER_CONTAINER_NAME], check=False)
    subprocess.run(
        [
            'docker',
            'build',
            '-t',
            SUPABASE_DOCKER_CONTAINER_NAME,
            '-f',
            str(SUPABASE_DOCKER_IMAGE_PATH),
            str(SUPABASE_DOCKER_IMAGE_PATH.parent),
        ],
        check=True,
    )
    subprocess.run(
        [
            'docker',
            'run',
            '-d',
            '--name',
            SUPABASE_DOCKER_CONTAINER_NAME,
            '-p',
            f'{TEST_SUPABASE_PORT}:5432',
            SUPABASE_DOCKER_CONTAINER_NAME,
        ],
        check=True,
    )


def supabase_stop_docker():
    """Stop the Docker container."""
    print("Stopping Supabase/Postgres Docker container...")
    subprocess.run(['docker', 'stop', SUPABASE_DOCKER_CONTAINER_NAME], check=True)


def supabase_verify_connection(conn):
    """Verify the connection to the database uses the test environment."""
    assert conn.info.host == TEST_SUPABASE_HOST
    assert conn.info.port == TEST_SUPABASE_PORT
    assert conn.info.dbname == TEST_SUPABASE_NAME
    assert conn.info.user == TEST_SUPABASE_USER
    assert conn.info.password == TEST_SUPABASE_PASSWORD


async def supabase_run_migrations(conn):
    """Run database schema initialization and migrations."""

    print("Running Supabase migrations...")
    with conn.cursor() as cur:
        # we need to do some stuff to get the new database to match what we expect
        # in Supabase. Some of these `SET`s might not be necessary, but they don't hurt.
        setup_query = """
        SET statement_timeout = 0;
        SET lock_timeout = 0;
        SET idle_in_transaction_session_timeout = 0;
        SET client_encoding = 'UTF8';
        SET standard_conforming_strings = on;
        SET search_path TO public;
        SET check_function_bodies = false;
        SET xmloption = content;
        SET client_min_messages = warning;
        SET row_security = off;

        -- Create essential extensions
        CREATE EXTENSION IF NOT EXISTS pg_stat_statements;
        CREATE EXTENSION IF NOT EXISTS "uuid-ossp";
        CREATE EXTENSION IF NOT EXISTS pgcrypto;

        -- Create auth schema and dummy JWT function
        CREATE OR REPLACE FUNCTION auth.jwt() RETURNS jsonb AS $$
        BEGIN
            RETURN '{"role": "authenticated", "sub": "00000000-0000-0000-0000-000000000000"}'::jsonb;
        END;
        $$ LANGUAGE plpgsql;

        -- the production bucket table differs from the one in the docker image
        -- so we drop it and recreate it
        DROP TABLE IF EXISTS storage.buckets CASCADE;
        CREATE TABLE IF NOT EXISTS storage.buckets (
            id text NOT NULL,
            name text NOT NULL,
            owner uuid,
            created_at timestamp with time zone DEFAULT now(),
            updated_at timestamp with time zone DEFAULT now(),
            public boolean DEFAULT false,
            avif_autodetection boolean DEFAULT false,
            file_size_limit bigint,
            allowed_mime_types text[],
            CONSTRAINT buckets_pkey PRIMARY KEY (id)
        );

        -- create two test users in the auth.users table so that we can reference them
        -- TODO move these to be inside the test_user, test_user2, and test_user3 fixtures
        -- so that it's obvious they are dependencies
        INSERT INTO auth.users (
            id,
            instance_id,
            aud,
            role,
            email,
            encrypted_password,
            confirmed_at,
            created_at,
            updated_at
        ) VALUES (
            '00000000-0000-0000-0000-000000000000',
            '00000000-0000-0000-0000-000000000000',
            'authenticated',
            'authenticated',
            'test@example.com',
            '$2a$10$abcdefghijklmnopqrstuvwxyz0123456789',
            NOW(),
            NOW(),
            NOW()
        ), (
            '00000000-0000-0000-0000-000000000001',
            '00000000-0000-0000-0000-000000000000',
            'authenticated',
            'authenticated',
            'test2@example.com',
            '$2a$10$abcdefghijklmnopqrstuvwxyz0123456789',
            NOW(),
            NOW(),
            NOW()
        ), (
            '00000000-0000-0000-0000-000000000002',
            '00000000-0000-0000-0000-000000000000',
            'authenticated',
            'authenticated',
            'test3@example.com',
            '$2a$10$abcdefghijklmnopqrstuvwxyz0123456789',
            NOW(),
            NOW(),
            NOW()
        );
        """
        cur.execute(setup_query)

        # migration files sorted by name
        migrations_dir = Path(__file__).parents[3] / 'supabase' / 'migrations'
        migration_files = sorted(migrations_dir.glob('*.sql'))

        print(f"Running {len(migration_files)} migrations...")

        for migration_file in migration_files:
            print(f"Running migration: {os.path.basename(migration_file)}")
            with open(migration_file) as f:
                sql = f.read()
                cur.execute(sql)

        # update public.orgs to have a `subscription_id TEXT DEFAULT ''` because this is in
        # prod, but not in the schema migrations.
        orgs_subscription_id_query = """
        ALTER TABLE public.orgs
        ADD COLUMN IF NOT EXISTS subscription_id TEXT DEFAULT '';
        """
        cur.execute(orgs_subscription_id_query)

    conn.commit()  # must commit transaction to apply changes
    print("Migrations complete!")


@pytest.fixture(scope="session", autouse=True)
async def supabase_setup_db_server():
    """Configure test database server and connection."""
    import time
    import psycopg
    from agentops.common.postgres import get_connection, close_connection, ConnectionConfig

    # Close any existing connections (sweaty)
    close_connection()

    if not is_github_actions():
        supabase_start_docker()

    # NOTE: we do not need to save or restore these variables as the tests will run
    # and then we'll be done
    ConnectionConfig.host = TEST_SUPABASE_HOST
    ConnectionConfig.port = TEST_SUPABASE_PORT
    ConnectionConfig.database = TEST_SUPABASE_NAME
    ConnectionConfig.user = TEST_SUPABASE_USER
    ConnectionConfig.password = TEST_SUPABASE_PASSWORD

    print(f"Waiting for PostgreSQL to be ready at {ConnectionConfig.host}:{ConnectionConfig.port}...")

    for attempt in range(30):
        try:
            # Use psycopg directly for the health check
            conn = psycopg.connect(ConnectionConfig.to_connection_string())
            with conn.cursor() as cur:
                cur.execute("SELECT 1")
            conn.close()
            print("PostgreSQL is ready!")
            break
        except Exception:
            print('.', end='', flush=True)
            time.sleep(2)
    else:
        raise Exception("Failed to connect to PostgreSQL after multiple attempts")

    pool = get_connection()
    with pool.connection() as conn:
        # double-check connection info (sweaty)
        supabase_verify_connection(conn)

        # run migrations once for the session
        await supabase_run_migrations(conn)

    yield

    if not is_github_actions():
        supabase_stop_docker()


@pytest.fixture(scope="session")
async def db_session():
    """Session fixture providing a database connection."""
    from agentops.common.postgres import get_connection, close_connection

    print("Creating database connection...")
    pool = get_connection()
    with pool.connection() as conn:
        try:
            # double-check connection uses test creds (sweaty)
            supabase_verify_connection(conn)

            yield conn
        except Exception as e:
            print(f"Error during database session: {str(e)}")
            conn.rollback()
        finally:
            close_connection()


@pytest.fixture(scope="session")
async def orm_session():
    """Fixture to provide an orm session for tests."""
    # does not automatically commit or rollback
    from agentops.common.orm import get_orm_session

    session = next(get_orm_session())

    # double-check connection uses test creds (sweaty)
    conn = session.connection()
    supabase_verify_connection(conn.connection.driver_connection)

    yield session

    session.close()