WonderQ-Project/WonderQ-Admin/tests/test_clean_database_migration_path.py

from importlib.util import module_from_spec, spec_from_file_location
import subprocess
from pathlib import Path
import sys

import pytest

VERSIONS_DIR = Path(__file__).parents[1] / "alembic" / "versions"


def load_migration(filename: str):
    path = VERSIONS_DIR / filename
    spec = spec_from_file_location(f"clean_database_{path.stem}", path)
    assert spec and spec.loader
    migration = module_from_spec(spec)
    spec.loader.exec_module(migration)
    return migration


class ExistingRows:
    def first(self):
        return ("existing",)


class EmptyRows:
    def first(self):
        return None


class ExistingBind:
    def execute(self, _statement):
        return ExistingRows()


class EmptyBind:
    def execute(self, _statement):
        return EmptyRows()


class CurrentSchemaInspector:
    tables = {
        "AdminDepartment",
        "AdminMenu",
        "AdminRole",
        "AdminRoleDepartment",
        "AdminRoleMenu",
        "AdminUser",
        "AdminUserDepartment",
        "AdminUserRole",
        "AuditLog",
        "ConciergeAdvisor",
        "Customer",
        "CustomerBrowseHistory",
        "DetailRecord",
        "HeroSlide",
        "HomeTeamBuilding",
        "HomeWanfaRecommendation",
        "HomeWildArchive",
        "Lead",
        "LeadFollowup",
        "MediaAsset",
        "Order",
        "VehicleOption",
        "WanfaCategory",
        "WanfaRoute",
    }

    def get_table_names(self):
        return sorted(self.tables)

    def get_columns(self, table_name):
        if table_name not in self.tables:
            return []
        return [{"name": "deptId"}, {"name": "createdById"}]

    def get_indexes(self, table_name):
        if table_name not in self.tables:
            return []
        return [
            {"name": f"ix_{table_name}_deptId"},
            {"name": f"ix_{table_name}_createdById"},
        ]

    def get_foreign_keys(self, table_name):
        if table_name not in self.tables:
            return []
        return [
            {
                "name": f"fk_{table_name}_deptId",
                "constrained_columns": ["deptId"],
                "referred_table": "AdminDepartment",
                "referred_columns": ["id"],
            },
            {
                "name": f"fk_{table_name}_createdById",
                "constrained_columns": ["createdById"],
                "referred_table": "AdminUser",
                "referred_columns": ["id"],
            },
        ]


def test_admin_rbac_does_not_recreate_tables_already_created_by_initial_baseline(monkeypatch):
    migration = load_migration("0025_admin_rbac.py")
    inspector = CurrentSchemaInspector()
    calls = []

    monkeypatch.setattr(migration, "is_offline_mode", lambda: False, raising=False)
    monkeypatch.setattr(migration, "inspect", lambda _bind: inspector, raising=False)
    monkeypatch.setattr(migration.op, "get_bind", lambda: ExistingBind())
    monkeypatch.setattr(
        migration.op,
        "create_table",
        lambda *args, **_kwargs: calls.append(("create_table", args[0])),
    )
    monkeypatch.setattr(
        migration.op,
        "bulk_insert",
        lambda *args, **_kwargs: calls.append(("bulk_insert", args[0].name)),
    )
    monkeypatch.setattr(
        migration.op,
        "execute",
        lambda *_args, **_kwargs: calls.append(("execute",)),
    )

    migration.upgrade()

    assert [call for call in calls if call[0] in {"create_table", "bulk_insert"}] == []


def test_initial_baseline_seeds_default_department_before_owned_data(monkeypatch):
    migration = load_migration("0001_initial_schema.py")
    inspector = CurrentSchemaInspector()
    calls = []

    monkeypatch.setattr(migration.context, "is_offline_mode", lambda: False)
    monkeypatch.setattr(migration, "inspect", lambda _bind: inspector, raising=False)
    monkeypatch.setattr(migration.op, "get_bind", lambda: EmptyBind())
    monkeypatch.setattr(
        migration.op,
        "bulk_insert",
        lambda table, rows: calls.append((table.name, rows)),
    )

    migration.upgrade()

    assert [table_name for table_name, _rows in calls] == ["AdminDepartment"]
    assert calls[0][1][0]["id"] == migration.DEFAULT_DEPT_ID


def test_concierge_details_migration_converts_json_to_jsonb(monkeypatch):
    migration_path = VERSIONS_DIR / "0035_concierge_details_jsonb.py"
    if not migration_path.exists():
        pytest.fail("0035_concierge_details_jsonb.py is required to align the model type")
    migration = load_migration("0035_concierge_details_jsonb.py")
    calls = []

    monkeypatch.setattr(
        migration.op,
        "alter_column",
        lambda *args, **kwargs: calls.append((args, kwargs)),
    )

    migration.upgrade()

    assert migration.down_revision == "0034_media_assets"
    assert calls[0][0][:2] == ("ConciergeAdvisor", "details")
    assert calls[0][1]["postgresql_using"] == '"details"::jsonb'


def test_offline_upgrade_generates_sql_for_a_clean_database():
    result = subprocess.run(
        [sys.executable, "-m", "alembic", "upgrade", "head", "--sql"],
        cwd=VERSIONS_DIR.parents[1],
        capture_output=True,
        text=True,
        timeout=30,
    )

    assert result.returncode == 0, result.stderr[-4000:]
    assert "0035_concierge_details_jsonb" in result.stdout
    assert result.stdout.count('CREATE TABLE "AdminRole"') == 1
    assert result.stdout.count('CREATE TABLE "MediaAsset"') == 1
    assert 'ALTER TABLE "HomeWildArchive" ADD COLUMN images' not in result.stdout
    assert 'ALTER TABLE "HomeTeamBuilding" ADD COLUMN "detailSubtitle"' not in result.stdout
    assert 'ALTER TABLE "DetailRecord" ADD COLUMN "conciergeAdvisorId"' not in result.stdout
    assert 'ALTER TABLE "Lead" ADD COLUMN "leadType"' not in result.stdout


def test_admin_ownership_skips_missing_and_already_owned_tables(monkeypatch):
    migration = load_migration("0026_admin_ownership.py")
    inspector = CurrentSchemaInspector()
    calls = []

    monkeypatch.setattr(migration, "is_offline_mode", lambda: False, raising=False)
    monkeypatch.setattr(migration, "inspect", lambda _bind: inspector, raising=False)
    monkeypatch.setattr(migration.op, "get_bind", lambda: ExistingBind())
    for operation in (
        "add_column",
        "create_index",
        "create_foreign_key",
        "execute",
        "alter_column",
    ):
        monkeypatch.setattr(
            migration.op,
            operation,
            lambda *args, _operation=operation, **_kwargs: calls.append(
                (_operation, args[0] if args else None)
            ),
        )

    migration.upgrade()

    assert calls == []