sqlalchemy-pgview/tests/conftest.py

"""Pytest configuration and fixtures."""

import pytest
from sqlalchemy import (
    Column,
    DateTime,
    ForeignKey,
    Integer,
    MetaData,
    Numeric,
    String,
    Table,
    create_engine,
    func,
    insert,
)
from sqlalchemy.engine import Engine


@pytest.fixture
def metadata() -> MetaData:
    """Create a fresh MetaData instance."""
    return MetaData()


@pytest.fixture
def pg_engine() -> Engine | None:
    """Create a PostgreSQL engine if available.

    Set the POSTGRES_URL environment variable to run PostgreSQL tests.
    Example: postgresql://user:pass@localhost:5432/testdb
    """
    import os

    url = os.environ.get("POSTGRES_URL")
    if not url:
        pytest.skip("POSTGRES_URL not set")
    return create_engine(url)


@pytest.fixture
def sample_tables(metadata: MetaData) -> tuple[Table, Table]:
    """Create sample tables for testing."""
    users = Table(
        "users",
        metadata,
        Column("id", Integer, primary_key=True),
        Column("name", String(100)),
        Column("email", String(100)),
    )

    orders = Table(
        "orders",
        metadata,
        Column("id", Integer, primary_key=True),
        Column("user_id", Integer),
        Column("total", Integer),
    )

    return users, orders


@pytest.fixture
def one_to_many_tables(metadata: MetaData) -> dict[str, Table]:
    """Create tables with one-to-many relationships.

    Schema:
        authors (1) --< (many) books
        books (1) --< (many) reviews
    """
    authors = Table(
        "authors",
        metadata,
        Column("id", Integer, primary_key=True),
        Column("name", String(100), nullable=False),
        Column("country", String(50)),
    )

    books = Table(
        "books",
        metadata,
        Column("id", Integer, primary_key=True),
        Column("title", String(200), nullable=False),
        Column("author_id", Integer, ForeignKey("authors.id"), nullable=False),
        Column("price", Numeric(10, 2)),
        Column("published_at", DateTime),
    )

    reviews = Table(
        "reviews",
        metadata,
        Column("id", Integer, primary_key=True),
        Column("book_id", Integer, ForeignKey("books.id"), nullable=False),
        Column("rating", Integer, nullable=False),
        Column("comment", String(500)),
    )

    return {"authors": authors, "books": books, "reviews": reviews}


@pytest.fixture
def many_to_many_tables(metadata: MetaData) -> dict[str, Table]:
    """Create tables with many-to-many relationships.

    Schema:
        students (many) --< student_courses >-- (many) courses
        courses (many) --< course_tags >-- (many) tags
    """
    students = Table(
        "students",
        metadata,
        Column("id", Integer, primary_key=True),
        Column("name", String(100), nullable=False),
        Column("email", String(100), unique=True),
    )

    courses = Table(
        "courses",
        metadata,
        Column("id", Integer, primary_key=True),
        Column("name", String(100), nullable=False),
        Column("credits", Integer, default=3),
    )

    student_courses = Table(
        "student_courses",
        metadata,
        Column("student_id", Integer, ForeignKey("students.id"), primary_key=True),
        Column("course_id", Integer, ForeignKey("courses.id"), primary_key=True),
        Column("grade", Numeric(3, 2)),
        Column("enrolled_at", DateTime, server_default=func.now()),
    )

    tags = Table(
        "tags",
        metadata,
        Column("id", Integer, primary_key=True),
        Column("name", String(50), unique=True, nullable=False),
    )

    course_tags = Table(
        "course_tags",
        metadata,
        Column("course_id", Integer, ForeignKey("courses.id"), primary_key=True),
        Column("tag_id", Integer, ForeignKey("tags.id"), primary_key=True),
    )

    return {
        "students": students,
        "courses": courses,
        "student_courses": student_courses,
        "tags": tags,
        "course_tags": course_tags,
    }


@pytest.fixture
def pg_one_to_many_tables(
    pg_engine: Engine, one_to_many_tables: dict[str, Table], metadata: MetaData
) -> dict[str, Table]:
    """Create one-to-many tables in PostgreSQL with sample data."""
    metadata.create_all(pg_engine)

    with pg_engine.begin() as conn:
        # Insert authors
        conn.execute(
            insert(one_to_many_tables["authors"]),
            [
                {"id": 1, "name": "George Orwell", "country": "UK"},
                {"id": 2, "name": "Gabriel Garcia Marquez", "country": "Colombia"},
                {"id": 3, "name": "Haruki Murakami", "country": "Japan"},
            ],
        )

        # Insert books
        conn.execute(
            insert(one_to_many_tables["books"]),
            [
                {"id": 1, "title": "1984", "author_id": 1, "price": 15.99},
                {"id": 2, "title": "Animal Farm", "author_id": 1, "price": 12.99},
                {"id": 3, "title": "One Hundred Years of Solitude", "author_id": 2, "price": 18.99},
                {"id": 4, "title": "Norwegian Wood", "author_id": 3, "price": 14.99},
                {"id": 5, "title": "Kafka on the Shore", "author_id": 3, "price": 16.99},
            ],
        )

        # Insert reviews
        conn.execute(
            insert(one_to_many_tables["reviews"]),
            [
                {"id": 1, "book_id": 1, "rating": 5, "comment": "A masterpiece"},
                {"id": 2, "book_id": 1, "rating": 4, "comment": "Thought-provoking"},
                {"id": 3, "book_id": 1, "rating": 5, "comment": "Must read"},
                {"id": 4, "book_id": 2, "rating": 4, "comment": "Great allegory"},
                {"id": 5, "book_id": 3, "rating": 5, "comment": "Beautiful prose"},
                {"id": 6, "book_id": 4, "rating": 4, "comment": "Melancholic"},
                {"id": 7, "book_id": 5, "rating": 5, "comment": "Surreal and captivating"},
            ],
        )

    yield one_to_many_tables

    metadata.drop_all(pg_engine)


@pytest.fixture
def pg_many_to_many_tables(
    pg_engine: Engine, many_to_many_tables: dict[str, Table], metadata: MetaData
) -> dict[str, Table]:
    """Create many-to-many tables in PostgreSQL with sample data."""
    metadata.create_all(pg_engine)

    with pg_engine.begin() as conn:
        # Insert students
        conn.execute(
            insert(many_to_many_tables["students"]),
            [
                {"id": 1, "name": "Alice", "email": "alice@example.com"},
                {"id": 2, "name": "Bob", "email": "bob@example.com"},
                {"id": 3, "name": "Charlie", "email": "charlie@example.com"},
            ],
        )

        # Insert courses
        conn.execute(
            insert(many_to_many_tables["courses"]),
            [
                {"id": 1, "name": "Database Systems", "credits": 4},
                {"id": 2, "name": "Web Development", "credits": 3},
                {"id": 3, "name": "Machine Learning", "credits": 4},
            ],
        )

        # Insert student_courses (enrollments)
        conn.execute(
            insert(many_to_many_tables["student_courses"]),
            [
                {"student_id": 1, "course_id": 1, "grade": 3.8},
                {"student_id": 1, "course_id": 2, "grade": 4.0},
                {"student_id": 1, "course_id": 3, "grade": 3.5},
                {"student_id": 2, "course_id": 1, "grade": 3.2},
                {"student_id": 2, "course_id": 3, "grade": 3.9},
                {"student_id": 3, "course_id": 2, "grade": 3.7},
            ],
        )

        # Insert tags
        conn.execute(
            insert(many_to_many_tables["tags"]),
            [
                {"id": 1, "name": "programming"},
                {"id": 2, "name": "data"},
                {"id": 3, "name": "ai"},
            ],
        )

        # Insert course_tags
        conn.execute(
            insert(many_to_many_tables["course_tags"]),
            [
                {"course_id": 1, "tag_id": 2},  # Database -> data
                {"course_id": 2, "tag_id": 1},  # Web Dev -> programming
                {"course_id": 3, "tag_id": 2},  # ML -> data
                {"course_id": 3, "tag_id": 3},  # ML -> ai
            ],
        )

    yield many_to_many_tables

    metadata.drop_all(pg_engine)