ZY-Agent/backend/tests/test_alembic_default_workspace_id.py

"""Alembic revision 0001 round-trips on both Postgres and SQLite.

Verifies the first DeerFlow migration adds `users.default_workspace_id`
(with FK to `workspaces`) on upgrade and removes it on downgrade. Both
backends are exercised because the migration relies on
`op.batch_alter_table` for SQLite ALTER compatibility — we want to know
if either dialect regresses.

These tests are synchronous: alembic's command layer is sync, and our
`env.py` calls `asyncio.run(...)` internally. Running under
pytest-anyio would put us inside an event loop and crash that
`asyncio.run` call, so we keep the test bodies plain `def`.
"""

from __future__ import annotations

import tempfile
from pathlib import Path

import pytest
from alembic import command
from alembic.config import Config
from sqlalchemy import create_engine, inspect, text

_ALEMBIC_INI = Path(__file__).resolve().parents[1] / "packages" / "harness" / "deerflow" / "persistence" / "migrations" / "alembic.ini"


def _make_alembic_config(url: str) -> Config:
    cfg = Config(str(_ALEMBIC_INI))
    cfg.set_main_option("sqlalchemy.url", url)
    return cfg


def _bootstrap_pre_pr4_schema(sync_url: str) -> None:
    """Create the minimal pre-PR4 schema the migration needs to ALTER.

    Only `users` (without `default_workspace_id`) and `workspaces` (just `id`)
    are required so the FK target resolves. The rest of the production
    schema is irrelevant to this migration.
    """
    engine = create_engine(sync_url)
    with engine.begin() as conn:
        conn.execute(
            text(
                """
                CREATE TABLE workspaces (
                    id VARCHAR(36) PRIMARY KEY,
                    name VARCHAR(64) NOT NULL
                )
                """
            )
        )
        conn.execute(
            text(
                """
                CREATE TABLE users (
                    id VARCHAR(36) PRIMARY KEY,
                    email VARCHAR(255) NOT NULL
                )
                """
            )
        )
    engine.dispose()


def _assert_column_present(sync_url: str) -> None:
    engine = create_engine(sync_url)
    insp = inspect(engine)
    cols = {c["name"] for c in insp.get_columns("users")}
    assert "default_workspace_id" in cols, f"column missing; got {cols}"
    fks = insp.get_foreign_keys("users")
    fk_to_ws = [fk for fk in fks if fk.get("referred_table") == "workspaces"]
    assert fk_to_ws, f"FK to workspaces missing; got {fks}"
    assert fk_to_ws[0]["constrained_columns"] == ["default_workspace_id"]
    assert fk_to_ws[0]["referred_columns"] == ["id"]
    engine.dispose()


def _assert_column_absent(sync_url: str) -> None:
    engine = create_engine(sync_url)
    insp = inspect(engine)
    cols = {c["name"] for c in insp.get_columns("users")}
    assert "default_workspace_id" not in cols, f"column still present; got {cols}"
    engine.dispose()


# ---------- SQLite tests ----------------------------------------------------


def test_sqlite_upgrade_adds_default_workspace_id_with_fk() -> None:
    """SQLite: upgrade 0001 adds column + FK; batch_alter_table works."""
    with tempfile.TemporaryDirectory() as tmp:
        db_path = Path(tmp) / "test.db"
        sync_url = f"sqlite:///{db_path}"
        async_url = f"sqlite+aiosqlite:///{db_path}"

        _bootstrap_pre_pr4_schema(sync_url)

        cfg = _make_alembic_config(async_url)
        command.upgrade(cfg, "0001_users_default_workspace")

        _assert_column_present(sync_url)


def test_sqlite_downgrade_removes_default_workspace_id() -> None:
    """SQLite: downgrade 0001 removes the column it added."""
    with tempfile.TemporaryDirectory() as tmp:
        db_path = Path(tmp) / "test.db"
        sync_url = f"sqlite:///{db_path}"
        async_url = f"sqlite+aiosqlite:///{db_path}"

        _bootstrap_pre_pr4_schema(sync_url)

        cfg = _make_alembic_config(async_url)
        command.upgrade(cfg, "0001_users_default_workspace")
        _assert_column_present(sync_url)

        command.downgrade(cfg, "-1")
        _assert_column_absent(sync_url)


# ---------- Postgres tests --------------------------------------------------


@pytest.mark.postgres
def test_postgres_upgrade_adds_default_workspace_id_with_fk(postgres_url: str) -> None:
    """Postgres: upgrade 0001 adds column + FK pointing at workspaces(id)."""
    sync_url = postgres_url.replace("+asyncpg", "+psycopg")
    _bootstrap_pre_pr4_schema(sync_url)

    cfg = _make_alembic_config(postgres_url)
    command.upgrade(cfg, "head")

    _assert_column_present(sync_url)


@pytest.mark.postgres
def test_postgres_downgrade_removes_default_workspace_id(postgres_url: str) -> None:
    """Postgres: downgrade 0001 cleanly drops the FK and column."""
    sync_url = postgres_url.replace("+asyncpg", "+psycopg")
    _bootstrap_pre_pr4_schema(sync_url)

    cfg = _make_alembic_config(postgres_url)
    command.upgrade(cfg, "head")
    _assert_column_present(sync_url)

    command.downgrade(cfg, "-1")
    _assert_column_absent(sync_url)