zhixing-system/zhixing-server/tests/integration/test_market_data_migration.py

import os
from pathlib import Path

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

from zhixing_server.bootstrap.config import get_settings, sqlalchemy_database_url


@pytest.mark.integration
def test_postgres_migration_creates_market_data_contract(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    database_url = os.getenv("ZHIXING_TEST_DATABASE_URL")
    if not database_url:
        pytest.skip("set ZHIXING_TEST_DATABASE_URL to run PostgreSQL integration tests")

    monkeypatch.setenv("ZHIXING_DATABASE_URL", database_url)
    get_settings.cache_clear()
    server_root = Path(__file__).parents[2]
    config = Config(str(server_root / "alembic.ini"))
    sqlalchemy_url = sqlalchemy_database_url(database_url)
    config.set_main_option("sqlalchemy.url", sqlalchemy_url.replace("%", "%%"))
    engine: Engine = create_engine(sqlalchemy_url)
    command.upgrade(config, "head")
    try:
        tables = set(inspect(engine).get_table_names())
        assert {
            "market_stock",
            "market_daily_bar",
            "market_daily_basic",
            "market_sync_batch",
            "market_sync_item",
            "market_integrity_check",
            "market_integrity_issue",
            "selection_run",
            "selection_run_item",
            "selection_signal",
            "sector_radar_source_snapshot",
            "sector_radar_membership",
            "sector_radar_stock_fact",
            "sector_radar_publication",
            "sector_radar_ranking",
        } <= tables
    finally:
        engine.dispose()
        get_settings.cache_clear()