tests/e2e_ui/fixtures/database_fixtures.py

"""
Worker-based database isolation fixtures for pytest-xdist parallel execution.

Supports both PostgreSQL and SQLite:
- PostgreSQL: Each worker gets a separate schema for isolation
- SQLite: Shared database with transaction rollback for isolation

Usage:
    def test_my_feature(db_session):
        # db_session is automatically isolated
        # Transactions are rolled back after test
        agent = Agent(name="test")
        db_session.add(agent)
        db_session.commit()
        # Data automatically cleaned up after test
"""

import os
import sys
import pytest
import sqlalchemy
from sqlalchemy import text
from sqlalchemy.orm import sessionmaker, Session
from typing import Generator

# Add backend to path for imports
backend_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
if backend_dir not in sys.path:
    sys.path.insert(0, backend_dir)

from core.models import Base


@pytest.fixture(scope="session")
def is_sqlite(get_engine: sqlalchemy.Engine) -> bool:
    """
    Detect if using SQLite database.

    Args:
        get_engine: SQLAlchemy engine

    Returns:
        True if SQLite, False if PostgreSQL
    """
    return get_engine.dialect.name == "sqlite"


@pytest.fixture(scope="session")
def worker_schema(worker_id: str) -> str:
    """
    Return worker-specific schema name (PostgreSQL only).

    Args:
        worker_id: pytest-xdist worker ID fixture (e.g., 'gw0', 'gw1', 'master')

    Returns:
        Schema name for this worker (e.g., 'test_schema_gw0')

    Examples:
        >>> worker_schema("gw0")
        'test_schema_gw0'
        >>> worker_schema("master")
        'test_schema_master'
    """
    # Handle master process (no xdist)
    if worker_id == "master":
        return "test_schema_gw0"
    return f"test_schema_{worker_id}"


@pytest.fixture(scope="session")
def get_engine() -> sqlalchemy.Engine:
    """
    Create SQLAlchemy engine for test database.

    Uses DATABASE_URL environment variable or defaults to PostgreSQL.

    Returns:
        SQLAlchemy engine instance

    Note:
        Engine is created once per test session and reused across tests.
    """
    database_url = os.getenv(
        "DATABASE_URL",
        "postgresql://atom:atom_test@localhost:5432/atom_test"
    )
    engine = sqlalchemy.create_engine(database_url, pool_pre_ping=True)
    return engine


@pytest.fixture(scope="session", autouse=True)
def create_worker_schema(
    worker_schema: str,
    get_engine: sqlalchemy.Engine,
    is_sqlite: bool,
    request: pytest.FixtureRequest
) -> Generator[str, None, None]:
    """
    Create worker-specific schema before test session and drop after (PostgreSQL only).

    For SQLite, this fixture yields without schema operations.

    Args:
        worker_schema: Schema name for this worker
        get_engine: SQLAlchemy engine
        is_sqlite: True if using SQLite
        request: Pytest request object

    Yields:
        Schema name (or placeholder for SQLite)

    Note:
        - PostgreSQL: Schema created with CREATE SCHEMA IF NOT EXISTS
        - PostgreSQL: Dropped with CASCADE after test session
        - SQLite: No schema operations (SQLite doesn't support schemas)
        - Autouse=True: runs automatically for all tests
        - Skipped for tests marked with no_browser
    """
    # Skip for unit tests marked with no_browser
    if request.node.get_closest_marker('no_browser'):
        yield worker_schema
        return

    schema = worker_schema

    # Skip schema operations for SQLite
    if is_sqlite:
        yield schema
        return

    # Create schema (PostgreSQL only)
    with get_engine.connect() as conn:
        conn.execute(text(f"CREATE SCHEMA IF NOT EXISTS {schema}"))
        conn.commit()

    yield schema

    # Teardown: drop schema after session (PostgreSQL only)
    with get_engine.connect() as conn:
        conn.execute(text(f"DROP SCHEMA IF EXISTS {schema} CASCADE"))
        conn.commit()


@pytest.fixture(scope="session")
def init_db(
    create_worker_schema: str,
    get_engine: sqlalchemy.Engine,
    is_sqlite: bool
) -> Generator[None, None, None]:
    """
    Initialize database tables in worker schema (PostgreSQL) or main database (SQLite).

    Args:
        create_worker_schema: Schema name from create_worker_schema fixture
        get_engine: SQLAlchemy engine
        is_sqlite: True if using SQLite

    Yields:
        None

    Note:
        - PostgreSQL: Uses schema_translate_map to direct tables to worker schema
        - SQLite: Creates tables in main database
        - Drops all tables after test session
    """
    schema = create_worker_schema

    if is_sqlite:
        # SQLite: Create tables directly in main database
        Base.metadata.create_all(get_engine)
        yield
        Base.metadata.drop_all(get_engine)
    else:
        # PostgreSQL: Create engine with schema translation
        engine = get_engine.execution_options(
            schema_translate_map={None: schema}
        )
        # Create all tables
        Base.metadata.create_all(engine)
        yield
        # Drop all tables
        Base.metadata.drop_all(engine)


@pytest.fixture(scope="function")
def db_session(
    worker_schema: str,
    get_engine: sqlalchemy.Engine,
    init_db: None,
    is_sqlite: bool
) -> Generator[Session, None, None]:
    """
    Create database session with transaction rollback for test isolation.

    Args:
        worker_schema: Schema name for this worker (PostgreSQL only)
        get_engine: SQLAlchemy engine
        init_db: Database initialization fixture (ensures tables exist)
        is_sqlite: True if using SQLite

    Yields:
        SQLAlchemy session with transaction rollback

    Note:
        - PostgreSQL: Sets search_path to worker schema
        - Begins transaction before test
        - Rolls back transaction after test (no data pollution)
        - Closes session after rollback
    """
    # Create engine with REPEATABLE READ isolation (PostgreSQL) or default (SQLite)
    if is_sqlite:
        engine = get_engine
    else:
        engine = get_engine.execution_options(
            isolation_level="REPEATABLE READ"
        )

    # Create session with autocommit disabled for transaction control
    Session = sessionmaker(bind=engine, autocommit=False, expire_on_commit=False)
    session = Session()

    # Set search path to worker schema (PostgreSQL only)
    # Do this in a separate transaction
    if not is_sqlite:
        session.execute(text(f"SET search_path TO {worker_schema}"))
        session.commit()  # Commit the search_path change immediately
        # Begin a new transaction for the test
        session.begin()
    else:
        # For SQLite, transaction is already started by default
        pass

    yield session

    # Rollback transaction after test
    session.rollback()
    session.close()


@pytest.fixture(scope="session")
def drop_worker_schema(
    worker_schema: str,
    get_engine: sqlalchemy.Engine,
    is_sqlite: bool
) -> None:
    """
    Drop worker-specific schema after test session (PostgreSQL only).

    For SQLite, this fixture does nothing (no schemas to drop).

    Args:
        worker_schema: Schema name for this worker
        get_engine: SQLAlchemy engine
        is_sqlite: True if using SQLite

    Note:
        - PostgreSQL: Drops schema with CASCADE
        - SQLite: No operation (no schemas)
        - This fixture is kept for explicit cleanup if needed
        - create_worker_schema already handles cleanup via CASCADE
        - Use this for manual schema reset during debugging
    """
    if is_sqlite:
        return

    with get_engine.connect() as conn:
        conn.execute(text(f"DROP SCHEMA IF EXISTS {worker_schema} CASCADE"))
        conn.commit()