backend/tests/test_migrations.py

"""
Unit tests for database schema migrations (_migrate_schema).

Each test creates a minimal in-memory SQLite database that intentionally
omits one or more columns (simulating an older schema), runs
_migrate_schema(), and then verifies the missing columns were added
without corrupting existing data or dropping any existing columns.
"""
import pytest
from sqlalchemy import create_engine, inspect, text
from sqlalchemy.orm import sessionmaker


# ── Helpers ───────────────────────────────────────────────────────────────────

def _make_engine():
    """Fresh in-memory SQLite engine — fully isolated per test."""
    return create_engine(
        "sqlite:///:memory:",
        connect_args={"check_same_thread": False},
    )


def _columns(engine, table: str) -> set:
    """Return the set of column names for a table."""
    return {c["name"] for c in inspect(engine).get_columns(table)}


def _run_migrate(engine):
    """Patch app.database.engine and call _migrate_schema()."""
    import app.database as db_module
    from sqlalchemy import inspect as sa_inspect

    original_engine = db_module.engine
    db_module.engine = engine

    # Also patch inspect so _migrate_schema uses our engine's inspector
    original_inspect = db_module.inspect
    db_module.inspect = lambda _: sa_inspect(engine)

    try:
        db_module._migrate_schema()
    finally:
        db_module.engine = original_engine
        db_module.inspect = original_inspect


def _create_minimal_users(engine):
    """Create users table with only the original columns (no migration columns)."""
    with engine.begin() as conn:
        conn.execute(text("""
            CREATE TABLE users (
                id CHAR(36) PRIMARY KEY,
                username VARCHAR(80) NOT NULL UNIQUE,
                email VARCHAR(120) NOT NULL UNIQUE,
                hashed_password VARCHAR(255) NOT NULL,
                is_admin BOOLEAN DEFAULT FALSE,
                created_at TIMESTAMP
            )
        """))


def _create_minimal_documents(engine):
    """Create documents table with only the original columns."""
    with engine.begin() as conn:
        conn.execute(text("""
            CREATE TABLE documents (
                id CHAR(36) PRIMARY KEY,
                user_id CHAR(36) NOT NULL,
                filename VARCHAR(255) NOT NULL,
                original_name VARCHAR(255) NOT NULL,
                file_size INTEGER DEFAULT 0,
                page_count INTEGER DEFAULT 0,
                chunk_count INTEGER DEFAULT 0,
                status VARCHAR(20) DEFAULT 'pending',
                error_message TEXT,
                uploaded_at TIMESTAMP
            )
        """))


def _create_minimal_api_keys(engine):
    """Create api_keys table with only the original columns."""
    with engine.begin() as conn:
        conn.execute(text("""
            CREATE TABLE api_keys (
                id CHAR(36) PRIMARY KEY,
                user_id CHAR(36) NOT NULL,
                key_prefix VARCHAR(20) NOT NULL,
                hashed_key VARCHAR(255) NOT NULL UNIQUE
            )
        """))


def _create_minimal_chat_messages(engine):
    """Create chat_messages table without the feedback column."""
    with engine.begin() as conn:
        conn.execute(text("""
            CREATE TABLE chat_messages (
                id CHAR(36) PRIMARY KEY,
                user_id CHAR(36) NOT NULL,
                document_id CHAR(36),
                session_id CHAR(36),
                role VARCHAR(20) NOT NULL,
                content TEXT NOT NULL,
                sources_json TEXT,
                created_at TIMESTAMP
            )
        """))


def _setup_all_minimal_tables(engine):
    """Create all tables with older schemas so the migration script doesn't crash."""
    _create_minimal_users(engine)
    _create_minimal_documents(engine)
    _create_minimal_api_keys(engine)
    _create_minimal_chat_messages(engine)


# ── users migrations ──────────────────────────────────────────────────────────

def test_migrate_adds_hf_token_to_users():
    engine = _make_engine()
    _setup_all_minimal_tables(engine)
    assert "hf_token" not in _columns(engine, "users")
    _run_migrate(engine)
    assert "hf_token" in _columns(engine, "users")


def test_migrate_adds_role_to_users():
    engine = _make_engine()
    _setup_all_minimal_tables(engine)
    assert "role" not in _columns(engine, "users")
    _run_migrate(engine)
    assert "role" in _columns(engine, "users")


def test_migrate_adds_google_refresh_token_to_users():
    engine = _make_engine()
    _setup_all_minimal_tables(engine)
    assert "google_refresh_token" not in _columns(engine, "users")
    _run_migrate(engine)
    assert "google_refresh_token" in _columns(engine, "users")


def test_migrate_adds_last_login_to_users():
    engine = _make_engine()
    _setup_all_minimal_tables(engine)
    _run_migrate(engine)
    assert "last_login" in _columns(engine, "users")


def test_migrate_adds_is_verified_to_users():
    engine = _make_engine()
    _setup_all_minimal_tables(engine)
    assert "is_verified" not in _columns(engine, "users")
    _run_migrate(engine)
    assert "is_verified" in _columns(engine, "users")


def test_migrate_adds_verification_token_hash_to_users():
    engine = _make_engine()
    _setup_all_minimal_tables(engine)
    _run_migrate(engine)
    assert "verification_token_hash" in _columns(engine, "users")


def test_migrate_adds_verification_token_created_at_to_users():
    engine = _make_engine()
    _setup_all_minimal_tables(engine)
    _run_migrate(engine)
    assert "verification_token_created_at" in _columns(engine, "users")


def test_migrate_preserves_existing_users_columns():
    """Migration must never drop or rename any original users column."""
    engine = _make_engine()
    _setup_all_minimal_tables(engine)
    original = _columns(engine, "users")
    _run_migrate(engine)
    after = _columns(engine, "users")
    assert original.issubset(after), f"Columns removed: {original - after}"


def test_migrate_preserves_existing_user_data():
    """Rows inserted before migration must still be readable afterwards."""
    engine = _make_engine()
    _setup_all_minimal_tables(engine)
    with engine.begin() as conn:
        conn.execute(text(
            "INSERT INTO users (id, username, email, hashed_password) "
            "VALUES ('u1', 'alice', 'alice@example.com', 'hash')"
        ))
    _run_migrate(engine)
    with engine.connect() as conn:
        row = conn.execute(
            text("SELECT username, email FROM users WHERE id = 'u1'")
        ).fetchone()
    assert row is not None
    assert row[0] == "alice"
    assert row[1] == "alice@example.com"


# ── documents migrations ──────────────────────────────────────────────────────

def test_migrate_adds_is_deleted_to_documents():
    engine = _make_engine()
    _setup_all_minimal_tables(engine)
    assert "is_deleted" not in _columns(engine, "documents")
    _run_migrate(engine)
    assert "is_deleted" in _columns(engine, "documents")


def test_migrate_adds_summary_to_documents():
    engine = _make_engine()
    _setup_all_minimal_tables(engine)
    assert "summary" not in _columns(engine, "documents")
    _run_migrate(engine)
    assert "summary" in _columns(engine, "documents")


def test_migrate_adds_chunk_size_and_overlap_to_documents():
    engine = _make_engine()
    _setup_all_minimal_tables(engine)
    _run_migrate(engine)
    cols = _columns(engine, "documents")
    assert "chunk_size" in cols
    assert "chunk_overlap" in cols


def test_migrate_adds_processing_progress_to_documents():
    engine = _make_engine()
    _setup_all_minimal_tables(engine)
    _run_migrate(engine)
    assert "processing_progress" in _columns(engine, "documents")


def test_migrate_adds_processing_stage_to_documents():
    engine = _make_engine()
    _setup_all_minimal_tables(engine)
    _run_migrate(engine)
    assert "processing_stage" in _columns(engine, "documents")


def test_migrate_adds_extracted_urls_to_documents():
    engine = _make_engine()
    _setup_all_minimal_tables(engine)
    assert "extracted_urls" not in _columns(engine, "documents")
    _run_migrate(engine)
    assert "extracted_urls" in _columns(engine, "documents")


def test_migrate_adds_completed_at_to_documents():
    engine = _make_engine()
    _setup_all_minimal_tables(engine)
    _run_migrate(engine)
    assert "completed_at" in _columns(engine, "documents")


def test_migrate_adds_drive_columns_to_documents():
    engine = _make_engine()
    _setup_all_minimal_tables(engine)
    _run_migrate(engine)
    cols = _columns(engine, "documents")
    assert "drive_file_id" in cols
    assert "drive_folder_id" in cols
    assert "drive_synced_at" in cols


def test_migrate_preserves_existing_documents_columns():
    engine = _make_engine()
    _setup_all_minimal_tables(engine)
    original = _columns(engine, "documents")
    _run_migrate(engine)
    after = _columns(engine, "documents")
    assert original.issubset(after), f"Columns removed: {original - after}"


def test_migrate_preserves_existing_document_data():
    engine = _make_engine()
    _setup_all_minimal_tables(engine)
    with engine.begin() as conn:
        conn.execute(text(
            "INSERT INTO documents "
            "(id, user_id, filename, original_name, file_size, status) "
            "VALUES ('d1', 'u1', 'file.pdf', 'file.pdf', 1024, 'ready')"
        ))
    _run_migrate(engine)
    with engine.connect() as conn:
        row = conn.execute(
            text("SELECT original_name, status FROM documents WHERE id = 'd1'")
        ).fetchone()
    assert row is not None
    assert row[0] == "file.pdf"
    assert row[1] == "ready"


# ── api_keys migrations ───────────────────────────────────────────────────────

def test_migrate_adds_name_to_api_keys():
    engine = _make_engine()
    _setup_all_minimal_tables(engine)
    assert "name" not in _columns(engine, "api_keys")
    _run_migrate(engine)
    assert "name" in _columns(engine, "api_keys")


def test_migrate_adds_is_active_to_api_keys():
    engine = _make_engine()
    _setup_all_minimal_tables(engine)
    assert "is_active" not in _columns(engine, "api_keys")
    _run_migrate(engine)
    assert "is_active" in _columns(engine, "api_keys")


def test_migrate_adds_last_used_at_to_api_keys():
    engine = _make_engine()
    _setup_all_minimal_tables(engine)
    _run_migrate(engine)
    assert "last_used_at" in _columns(engine, "api_keys")


def test_migrate_preserves_existing_api_keys_columns():
    engine = _make_engine()
    _setup_all_minimal_tables(engine)
    original = _columns(engine, "api_keys")
    _run_migrate(engine)
    after = _columns(engine, "api_keys")
    assert original.issubset(after), f"Columns removed: {original - after}"


# ── chat_messages migrations ──────────────────────────────────────────────────

def test_migrate_adds_feedback_to_chat_messages():
    engine = _make_engine()
    _setup_all_minimal_tables(engine)
    assert "feedback" not in _columns(engine, "chat_messages")
    _run_migrate(engine)
    assert "feedback" in _columns(engine, "chat_messages")


def test_migrate_preserves_existing_chat_messages_columns():
    engine = _make_engine()
    _setup_all_minimal_tables(engine)
    original = _columns(engine, "chat_messages")
    _run_migrate(engine)
    after = _columns(engine, "chat_messages")
    assert original.issubset(after), f"Columns removed: {original - after}"


# ── Idempotency ───────────────────────────────────────────────────────────────

def test_migrate_is_idempotent_users():
    """Running _migrate_schema() twice must not raise or duplicate columns."""
    engine = _make_engine()
    _setup_all_minimal_tables(engine)
    _run_migrate(engine)
    cols_after_first = _columns(engine, "users")
    _run_migrate(engine)
    cols_after_second = _columns(engine, "users")
    assert cols_after_first == cols_after_second


def test_migrate_is_idempotent_documents():
    engine = _make_engine()
    _setup_all_minimal_tables(engine)
    _run_migrate(engine)
    cols_after_first = _columns(engine, "documents")
    _run_migrate(engine)
    cols_after_second = _columns(engine, "documents")
    assert cols_after_first == cols_after_second


def test_migrate_is_idempotent_api_keys():
    engine = _make_engine()
    _setup_all_minimal_tables(engine)
    _run_migrate(engine)
    cols_after_first = _columns(engine, "api_keys")
    _run_migrate(engine)
    cols_after_second = _columns(engine, "api_keys")
    assert cols_after_first == cols_after_second


def test_migrate_is_idempotent_chat_messages():
    engine = _make_engine()
    _setup_all_minimal_tables(engine)
    _run_migrate(engine)
    cols_after_first = _columns(engine, "chat_messages")
    _run_migrate(engine)
    cols_after_second = _columns(engine, "chat_messages")
    assert cols_after_first == cols_after_second


# ── Full schema (already migrated) ───────────────────────────────────────────

def test_migrate_on_fully_migrated_schema_is_safe():
    """If all columns already exist, _migrate_schema() must complete without error."""
    engine = _make_engine()
    _setup_all_minimal_tables(engine)
    # First pass adds all missing columns
    _run_migrate(engine)
    # Second pass should run safely as a complete no-op
    _run_migrate(engine)
    # Verify nothing was lost
    assert "hf_token" in _columns(engine, "users")
    assert "extracted_urls" in _columns(engine, "documents")
    assert "is_active" in _columns(engine, "api_keys")
    assert "feedback" in _columns(engine, "chat_messages")