ZY-Agent/backend/tests/test_backfill_workspace_id.py

"""Tests for ``scripts/backfill_workspace_id.py`` (Stage 0 PR5).

Each step in the three-step backfill is exercised in isolation against
a SQLite-on-disk database. The script wires in ``app.gateway.auth.workspace_slug``,
so we get the same slug semantics that the registration flow uses.

Pattern mirrors :mod:`test_workspace_repo`: ``init_engine`` + per-test
tmp_path, with explicit ``close_engine`` teardown so the singleton
session factory does not leak across tests.
"""

from __future__ import annotations

import uuid

import pytest
from sqlalchemy import select

from deerflow.persistence.feedback.model import FeedbackRow
from deerflow.persistence.models.run_event import RunEventRow
from deerflow.persistence.run.model import RunRow
from deerflow.persistence.thread_meta.model import ThreadMetaRow
from deerflow.persistence.user.model import UserRow
from deerflow.persistence.workspace.model import WorkspaceRow
from deerflow.persistence.workspace_membership.model import WorkspaceMembershipRow
from scripts.backfill_workspace_id import (
    LEGACY_WORKSPACE_ID,
    _ensure_legacy_workspace,
    _step1_create_workspaces_for_users,
    _step2_update_table_from_users,
    _step3_assign_legacy_workspace,
    backfill,
)

pytestmark = pytest.mark.anyio


@pytest.fixture
def anyio_backend() -> str:
    return "asyncio"


async def _init_engine(tmp_path):
    from deerflow.persistence.engine import get_session_factory, init_engine

    url = f"sqlite+aiosqlite:///{tmp_path / 'test.db'}"
    await init_engine("sqlite", url=url, sqlite_dir=str(tmp_path))
    return get_session_factory()


async def _close():
    from deerflow.persistence.engine import close_engine

    await close_engine()


async def _seed_user(sf, *, email: str, default_workspace_id: str | None = None, system_role: str = "user") -> str:
    user_id = str(uuid.uuid4())
    async with sf() as session:
        session.add(UserRow(id=user_id, email=email, default_workspace_id=default_workspace_id, system_role=system_role))
        await session.commit()
    return user_id


async def _seed_business_row(sf, model, **fields) -> None:
    async with sf() as session:
        session.add(model(**fields))
        await session.commit()


# ---------------------------------------------------------------------------
# Step 1: per-user workspace creation
# ---------------------------------------------------------------------------


async def test_creates_workspace_per_user_without_default(tmp_path):
    """Each user with NULL default_workspace_id gets a workspace + owner membership."""
    sf = await _init_engine(tmp_path)
    try:
        u_alice = await _seed_user(sf, email="alice@example.com")
        u_bob = await _seed_user(sf, email="bob+spam@example.com")

        count = await _step1_create_workspaces_for_users(sf, dry_run=False)
        assert count == 2

        async with sf() as session:
            workspaces = (await session.execute(select(WorkspaceRow))).scalars().all()
            memberships = (await session.execute(select(WorkspaceMembershipRow))).scalars().all()
            users = {u.id: u for u in (await session.execute(select(UserRow))).scalars().all()}

        # 2 workspaces, each with exactly one owner membership matching its user.
        assert len(workspaces) == 2
        assert len(memberships) == 2
        owners_by_ws = {m.workspace_id: m.user_id for m in memberships if m.role == "owner"}
        assert {m.role for m in memberships} == {"owner"}
        for ws in workspaces:
            assert owners_by_ws[ws.id] == ws.owner_id
            assert users[ws.owner_id].default_workspace_id == ws.id

        # Slug semantics: alice@ → "alice", bob+spam@ → "bob-spam".
        slugs = {ws.slug for ws in workspaces}
        assert slugs == {"alice", "bob-spam"}
        _ = u_alice, u_bob  # captured for readability
    finally:
        await _close()


async def test_step1_is_idempotent(tmp_path):
    """Second run is a no-op when every user already has a default_workspace_id."""
    sf = await _init_engine(tmp_path)
    try:
        await _seed_user(sf, email="carol@example.com")
        first = await _step1_create_workspaces_for_users(sf, dry_run=False)
        assert first == 1
        # Re-running picks up the just-populated default_workspace_id, so the
        # candidate set is empty.
        second = await _step1_create_workspaces_for_users(sf, dry_run=False)
        assert second == 0
        async with sf() as session:
            ws_count = len((await session.execute(select(WorkspaceRow))).scalars().all())
            mem_count = len((await session.execute(select(WorkspaceMembershipRow))).scalars().all())
        assert ws_count == 1
        assert mem_count == 1
    finally:
        await _close()


async def test_backfill_updates_4_tables_from_users(tmp_path):
    """Step 2 propagates each user's default_workspace_id into 4 business tables."""
    sf = await _init_engine(tmp_path)
    try:
        user_id = await _seed_user(sf, email="dave@example.com")
        # Pre-seed business rows owned by the user with NULL workspace_id.
        await _seed_business_row(sf, ThreadMetaRow, thread_id="t-1", user_id=user_id)
        await _seed_business_row(sf, RunRow, run_id="r-1", thread_id="t-1", user_id=user_id)
        await _seed_business_row(sf, FeedbackRow, feedback_id="f-1", thread_id="t-1", run_id="r-1", user_id=user_id, rating=1)
        await _seed_business_row(sf, RunEventRow, thread_id="t-1", run_id="r-1", user_id=user_id, event_type="lifecycle_started", category="lifecycle", seq=1)

        # Step 1 first so users.default_workspace_id is populated.
        await _step1_create_workspaces_for_users(sf, dry_run=False)

        async with sf() as session:
            ws_id = (await session.execute(select(WorkspaceRow.id))).scalar_one()

        # Step 2 updates each table.
        for table in ("threads_meta", "runs", "feedback", "run_events"):
            count = await _step2_update_table_from_users(sf, table, dry_run=False)
            assert count == 1, table

        async with sf() as session:
            tm = (await session.execute(select(ThreadMetaRow))).scalar_one()
            run = (await session.execute(select(RunRow))).scalar_one()
            fb = (await session.execute(select(FeedbackRow))).scalar_one()
            ev = (await session.execute(select(RunEventRow))).scalar_one()
        assert tm.workspace_id == ws_id
        assert run.workspace_id == ws_id
        assert fb.workspace_id == ws_id
        assert ev.workspace_id == ws_id

        # Re-running Step 2 is a no-op (filtered by workspace_id IS NULL).
        for table in ("threads_meta", "runs", "feedback", "run_events"):
            assert await _step2_update_table_from_users(sf, table, dry_run=False) == 0
    finally:
        await _close()


async def test_backfill_step2_isolates_per_user(tmp_path):
    """Two users with different default workspaces get their own threads tagged independently."""
    sf = await _init_engine(tmp_path)
    try:
        u_eve = await _seed_user(sf, email="eve@example.com")
        u_frank = await _seed_user(sf, email="frank@example.com")
        await _seed_business_row(sf, ThreadMetaRow, thread_id="t-eve", user_id=u_eve)
        await _seed_business_row(sf, ThreadMetaRow, thread_id="t-frank", user_id=u_frank)

        await _step1_create_workspaces_for_users(sf, dry_run=False)
        await _step2_update_table_from_users(sf, "threads_meta", dry_run=False)

        async with sf() as session:
            rows = {r.thread_id: r.workspace_id for r in (await session.execute(select(ThreadMetaRow))).scalars().all()}
            users = {u.id: u.default_workspace_id for u in (await session.execute(select(UserRow))).scalars().all()}
        assert rows["t-eve"] == users[u_eve]
        assert rows["t-frank"] == users[u_frank]
        assert rows["t-eve"] != rows["t-frank"]
    finally:
        await _close()


async def test_step1_skips_blacklisted_base_slug(tmp_path):
    """A user with email like admin@... gets bumped past the slug blacklist via the walker."""
    sf = await _init_engine(tmp_path)
    try:
        await _seed_user(sf, email="admin@example.com")
        await _step1_create_workspaces_for_users(sf, dry_run=False)
        async with sf() as session:
            ws = (await session.execute(select(WorkspaceRow))).scalar_one()
        # The walker treats "admin" as taken (blacklisted), so it falls
        # through to "admin-2" — the same behaviour the registration flow
        # uses for reserved slugs.
        assert ws.slug == "admin-2"
    finally:
        await _close()


# ---------------------------------------------------------------------------
# Step 3: orphan rows -> legacy_workspace
# ---------------------------------------------------------------------------


async def test_backfill_orphan_rows_go_to_legacy_workspace(tmp_path):
    """Rows with user_id=NULL get assigned the legacy_workspace UUID after Step 3."""
    sf = await _init_engine(tmp_path)
    try:
        # Seed a platform admin so the legacy workspace has an owner.
        await _seed_user(sf, email="admin@example.com", system_role="admin")
        # Orphan business rows (user_id=NULL): legacy data from before auth.
        await _seed_business_row(sf, ThreadMetaRow, thread_id="t-orphan", user_id=None)
        await _seed_business_row(sf, RunRow, run_id="r-orphan", thread_id="t-orphan", user_id=None)
        await _seed_business_row(sf, FeedbackRow, feedback_id="f-orphan", thread_id="t-orphan", run_id="r-orphan", user_id=None, rating=1)
        await _seed_business_row(sf, RunEventRow, thread_id="t-orphan", run_id="r-orphan", user_id=None, event_type="legacy", category="lifecycle", seq=1)

        # Ensure the anchor + reassign per table.
        created = await _ensure_legacy_workspace(sf, dry_run=False)
        assert created is True
        for table in ("threads_meta", "runs", "feedback", "run_events"):
            count = await _step3_assign_legacy_workspace(sf, table, dry_run=False)
            assert count == 1, table

        # Re-running the anchor helper is a no-op.
        assert await _ensure_legacy_workspace(sf, dry_run=False) is False

        async with sf() as session:
            tm = (await session.execute(select(ThreadMetaRow))).scalar_one()
            run = (await session.execute(select(RunRow))).scalar_one()
            fb = (await session.execute(select(FeedbackRow))).scalar_one()
            ev = (await session.execute(select(RunEventRow))).scalar_one()
            legacy = (await session.execute(select(WorkspaceRow).where(WorkspaceRow.id == LEGACY_WORKSPACE_ID))).scalar_one()
            legacy_mem = (await session.execute(select(WorkspaceMembershipRow).where(WorkspaceMembershipRow.workspace_id == LEGACY_WORKSPACE_ID))).scalar_one()

        assert tm.workspace_id == LEGACY_WORKSPACE_ID
        assert run.workspace_id == LEGACY_WORKSPACE_ID
        assert fb.workspace_id == LEGACY_WORKSPACE_ID
        assert ev.workspace_id == LEGACY_WORKSPACE_ID
        assert legacy.slug == "legacy"
        assert legacy_mem.role == "owner"
    finally:
        await _close()


async def test_ensure_legacy_workspace_refuses_when_no_users(tmp_path):
    """ensure_legacy_workspace raises a clear error if the DB has no users."""
    sf = await _init_engine(tmp_path)
    try:
        with pytest.raises(RuntimeError, match="no users exist"):
            await _ensure_legacy_workspace(sf, dry_run=False)
    finally:
        await _close()


async def test_backfill_dry_run_does_not_write(tmp_path):
    """``backfill(..., dry_run=True)`` reports counts but writes nothing."""
    sf = await _init_engine(tmp_path)
    try:
        await _seed_user(sf, email="admin@example.com", system_role="admin")
        user_id = await _seed_user(sf, email="helen@example.com")
        await _seed_business_row(sf, ThreadMetaRow, thread_id="t-owned", user_id=user_id)
        await _seed_business_row(sf, ThreadMetaRow, thread_id="t-orphan", user_id=None)
        await _seed_business_row(sf, RunRow, run_id="r-owned", thread_id="t-owned", user_id=user_id)

        # Snapshot row counts BEFORE the dry run so we can confirm
        # nothing changed AFTER.
        async with sf() as session:
            ws_before = len((await session.execute(select(WorkspaceRow))).scalars().all())
            mem_before = len((await session.execute(select(WorkspaceMembershipRow))).scalars().all())
            users_with_default_before = len((await session.execute(select(UserRow).where(UserRow.default_workspace_id.is_not(None)))).scalars().all())

        report = await backfill(sf, dry_run=True)
        assert report["dry_run"] is True
        # Step 1 reports 2 candidates (admin + helen, both without default).
        assert report["users_workspaces_created"] == 2
        # Step 2 reports 0 because Step 1 didn't actually populate
        # users.default_workspace_id under dry_run — the JOIN comes up empty.
        assert report["threads_meta_from_users"] == 0
        assert report["runs_from_users"] == 0
        # Step 3 reports the 3 NULL business rows (t-owned, t-orphan, r-owned).
        assert report["legacy_workspace_created"] is True
        assert report["threads_meta_legacy"] == 2
        assert report["runs_legacy"] == 1

        # State did not change.
        async with sf() as session:
            ws_after = len((await session.execute(select(WorkspaceRow))).scalars().all())
            mem_after = len((await session.execute(select(WorkspaceMembershipRow))).scalars().all())
            users_with_default_after = len((await session.execute(select(UserRow).where(UserRow.default_workspace_id.is_not(None)))).scalars().all())
            rows = (await session.execute(select(ThreadMetaRow.workspace_id))).scalars().all()
        assert ws_after == ws_before
        assert mem_after == mem_before
        assert users_with_default_after == users_with_default_before
        assert all(w is None for w in rows)
    finally:
        await _close()


async def test_full_backfill_orchestrator(tmp_path):
    """End-to-end: backfill() runs all three steps and reports per-step counts."""
    sf = await _init_engine(tmp_path)
    try:
        await _seed_user(sf, email="admin@example.com", system_role="admin")
        user_id = await _seed_user(sf, email="gina@example.com")
        await _seed_business_row(sf, ThreadMetaRow, thread_id="t-owned", user_id=user_id)
        await _seed_business_row(sf, ThreadMetaRow, thread_id="t-orphan", user_id=None)

        report = await backfill(sf, dry_run=False)
        assert report["dry_run"] is False
        # Two users were missing a default workspace (admin too — we
        # didn't pre-populate admin's default_workspace_id).
        assert report["users_workspaces_created"] == 2
        assert report["threads_meta_from_users"] == 1
        assert report["legacy_workspace_created"] is True
        assert report["threads_meta_legacy"] == 1

        async with sf() as session:
            rows = {r.thread_id: r.workspace_id for r in (await session.execute(select(ThreadMetaRow))).scalars().all()}
        assert rows["t-orphan"] == LEGACY_WORKSPACE_ID
        assert rows["t-owned"] != LEGACY_WORKSPACE_ID
        assert rows["t-owned"] is not None
    finally:
        await _close()