DataOps-platform/tests/integration/test_rule_publication_migration_upgrade.py

from __future__ import annotations

from contextlib import suppress

import pytest
from sqlalchemy import create_engine, text

from app.core.common.identifiers import new_governance_uid
from tests.integration.test_rule_artifact_migration_upgrade import (
    _compose_value,
    _upgrade,
)


@pytest.mark.parametrize("legacy_status", ["draft", "validated", "published"])
def test_any_legacy_rule_blocks_real_190_to_200_upgrade(legacy_status):
    platform_user = _compose_value(
        r"\n  postgres:.*?POSTGRES_USER:\s*([^\s]+)"
    )
    platform_password = _compose_value(
        r"\n  postgres:.*?POSTGRES_PASSWORD:\s*([^\s]+)"
    )
    platform_port = _compose_value(r'"(15432):5432"')
    admin_url = (
        f"postgresql+psycopg2://{platform_user}:{platform_password}"
        f"@127.0.0.1:{platform_port}/postgres"
    )
    database_name = (
        f"task7_legacy_{new_governance_uid().replace('-', '')}"
    )
    database_url = (
        f"postgresql+psycopg2://{platform_user}:{platform_password}"
        f"@127.0.0.1:{platform_port}/{database_name}"
    )
    admin = create_engine(admin_url, isolation_level="AUTOCOMMIT")
    engine = None
    try:
        with admin.connect() as connection:
            connection.execute(text(f'CREATE DATABASE "{database_name}"'))
        _upgrade(database_url, "20260723_190")
        engine = create_engine(database_url, pool_pre_ping=True)
        rule_uid = new_governance_uid()
        with engine.begin() as connection:
            connection.execute(
                text(
                    "INSERT INTO public.data_rules "
                    "(id, rule_uid, name, category, status) VALUES "
                    "(CAST(:id AS uuid), CAST(:rule_uid AS uuid), "
                    "'legacy rule', 'legacy', 'active')"
                ),
                {"id": new_governance_uid(), "rule_uid": rule_uid},
            )
            connection.execute(
                text(
                    "INSERT INTO public.data_rule_versions "
                    "(id, rule_uid, version_no, source_text, rule_spec, "
                    "spec_hash, status) VALUES "
                    "(CAST(:id AS uuid), CAST(:rule_uid AS uuid), 1, "
                    "'legacy', '{}'::jsonb, :spec_hash, :legacy_status)"
                ),
                {
                    "id": new_governance_uid(),
                    "rule_uid": rule_uid,
                    "spec_hash": "a" * 64,
                    "legacy_status": legacy_status,
                },
            )
        engine.dispose()
        engine = None

        with pytest.raises(Exception, match="pre-Task7"):
            _upgrade(database_url, "20260723_200")

        engine = create_engine(database_url, pool_pre_ping=True)
        with engine.connect() as connection:
            assert connection.execute(
                text("SELECT version_num FROM alembic_version")
            ).scalar_one() == "20260723_190"
    finally:
        if engine is not None:
            engine.dispose()
        with suppress(Exception), admin.connect() as connection:
            connection.execute(
                text(
                    "SELECT pg_terminate_backend(pid) "
                    "FROM pg_stat_activity "
                    "WHERE datname = :database_name "
                    "AND pid <> pg_backend_pid()"
                ),
                {"database_name": database_name},
            )
            connection.execute(
                text(f'DROP DATABASE IF EXISTS "{database_name}"')
            )
        admin.dispose()