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:
        inspector = inspect(engine)
        tables = set(inspector.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",
            "sector_radar_daily_aggregate",
            "sector_radar_publication_source",
        } <= tables
        item_columns = {column["name"] for column in inspector.get_columns("selection_run_item")}
        assert {
            "score_status",
            "score_value",
            "score_threshold",
            "score_version",
            "match_case_id",
            "match_case_name",
            "match_case_breakout_date",
            "match_breakdown",
            "score_reason",
        } <= item_columns
        constraint_names = {
            constraint["name"]
            for constraint in inspector.get_check_constraints("selection_run_item")
        }
        assert "ck_selection_run_item_breakdown_range" in constraint_names
        index_names = {index["name"] for index in inspector.get_indexes("selection_run_item")}
        assert "ix_selection_run_item_score" in index_names
    finally:
        engine.dispose()
        get_settings.cache_clear()