DRRK/tests/unit/test_migrations.py

from __future__ import annotations

import io
from pathlib import Path

import pytest
from sqlalchemy import Identity, create_engine, inspect, text

from alembic import command
from alembic.autogenerate import compare_metadata
from alembic.config import Config
from alembic.migration import MigrationContext
from app.core.config import settings
from app.db.base import Base

ROOT = Path(__file__).resolve().parents[2]


def _alembic_config(output_buffer: io.StringIO | None = None) -> Config:
    # Без alembic.ini: fileConfig() в env.py перенастроил бы логирование
    # всего процесса и сломал caplog в соседних тестах.
    cfg = Config(output_buffer=output_buffer)
    cfg.set_main_option("script_location", str(ROOT / "alembic"))
    return cfg


@pytest.fixture
def sqlite_url(tmp_path: Path, monkeypatch) -> str:
    db_path = (tmp_path / "migrations.db").as_posix()
    monkeypatch.setattr(settings, "DATABASE_URL", f"sqlite+aiosqlite:///{db_path}")
    monkeypatch.setattr(settings, "DATABASE_SCHEMA", "")
    return f"sqlite:///{db_path}"


def test_upgrade_creates_tables_and_seeds(sqlite_url: str) -> None:
    command.upgrade(_alembic_config(), "head")

    engine = create_engine(sqlite_url)
    with engine.connect() as conn:
        tables = set(inspect(conn).get_table_names())
        currencies = conn.execute(text("SELECT code FROM currency")).scalars().all()
        portfolios = conn.execute(text("SELECT code FROM portfolio")).scalars().all()
        op_types = conn.execute(
            text("SELECT code, created_by FROM operation_type")
        ).all()
    engine.dispose()

    assert {
        "currency",
        "portfolio",
        "operation_type",
        "security",
        "market_data",
        "exchange_rate",
        "operation",
        "external_request_log",
        "import_batch",
        "alembic_version",
    } <= tables
    assert sorted(currencies) == ["CNY", "RUB", "USD"]
    assert sorted(portfolios) == ["SALES", "UW"]
    assert sorted(code for code, _ in op_types) == [
        "ADJ",
        "BUY",
        "CPN",
        "NOM",
        "SELL",
        "UWFEE",
        "VAL",
    ]
    assert {created_by for _, created_by in op_types} == {"system"}


def test_migration_matches_models(sqlite_url: str) -> None:
    """Ревизия написана руками — сверяем её с моделями через autogenerate."""
    command.upgrade(_alembic_config(), "head")

    engine = create_engine(sqlite_url)
    with engine.connect() as conn:
        diff = compare_metadata(MigrationContext.configure(conn), Base.metadata)
    engine.dispose()

    assert [change for change in diff if not _is_sqlite_identity_noise(change)] == []


def _is_sqlite_identity_noise(change) -> bool:  # noqa: ANN001
    """SQLite не хранит IDENTITY, и при рефлексии его нет — autogenerate видит
    «modify_default id: None -> Identity()» у каждой таблицы. Это не расхождение
    ревизии с моделями (в Postgres-DDL identity есть, см. offline-тест ниже)."""
    if not isinstance(change, list) or len(change) != 1:
        return False
    kind, _, _, column, _, _, new = change[0]
    return kind == "modify_default" and column == "id" and isinstance(new, Identity)


def test_operation_save_revision_up_down(sqlite_url: str) -> None:
    """7c2e5a9d1f30 добавляет колонки/ограничения operation и типы UWFEE, ADJ;
    откат возвращает как было. Частичный индекс по trade_number переживает
    пересоздание таблицы в SQLite (batch) с тем же WHERE."""
    cfg = _alembic_config()
    command.upgrade(cfg, "head")
    engine = create_engine(sqlite_url)

    def snapshot() -> tuple[set[str], list[str], str]:
        with engine.connect() as conn:
            columns = {c["name"] for c in inspect(conn).get_columns("operation")}
            codes = conn.execute(text("SELECT code FROM operation_type")).scalars()
            index_sql = conn.execute(
                text(
                    "SELECT sql FROM sqlite_master "
                    "WHERE name = 'ux_operation_portfolio_trade_number'"
                )
            ).scalar_one()
        return columns, sorted(codes), index_sql

    columns, codes, index_sql = snapshot()
    assert {"currency_id", "settlement_amount", "source_row"} <= columns
    assert {"UWFEE", "ADJ"} <= set(codes)
    assert "WHERE trade_number IS NOT NULL" in index_sql

    command.downgrade(cfg, "3b9f1c2a7d41")
    columns, codes, index_sql = snapshot()
    assert not {"currency_id", "settlement_amount", "source_row"} & columns
    assert codes == ["BUY", "CPN", "NOM", "SELL", "VAL"]
    assert "WHERE trade_number IS NOT NULL" in index_sql

    command.upgrade(cfg, "head")
    assert snapshot()[1] == ["ADJ", "BUY", "CPN", "NOM", "SELL", "UWFEE", "VAL"]
    engine.dispose()


def test_downgrade_removes_everything(sqlite_url: str) -> None:
    command.upgrade(_alembic_config(), "head")
    command.downgrade(_alembic_config(), "base")

    engine = create_engine(sqlite_url)
    with engine.connect() as conn:
        tables = set(inspect(conn).get_table_names())
    engine.dispose()

    assert tables <= {"alembic_version"}


def test_offline_postgres_sql_uses_schema_jsonb_and_identity(monkeypatch) -> None:
    monkeypatch.setattr(
        settings, "DATABASE_URL", "postgresql+asyncpg://user:pass@localhost/drrk"
    )
    monkeypatch.setattr(settings, "DATABASE_SCHEMA", "drrk")
    buffer = io.StringIO()

    command.upgrade(_alembic_config(buffer), "head", sql=True)
    sql = buffer.getvalue()

    assert "CREATE SCHEMA IF NOT EXISTS drrk" not in sql  # только online-режим
    assert "CREATE TABLE drrk.market_data" in sql
    assert "raw JSONB" in sql
    assert "id BIGINT GENERATED BY DEFAULT AS IDENTITY" in sql
    assert "REFERENCES drrk.currency (id)" in sql
    assert "WHERE trade_number IS NOT NULL" in sql
    assert "CREATE TABLE drrk.alembic_version" in sql
    assert "CREATE TABLE drrk.import_batch" in sql
    assert "REFERENCES drrk.portfolio (id)" in sql
    assert "INSERT INTO drrk.currency" in sql
    assert "ALTER TABLE drrk.operation ADD COLUMN currency_id BIGINT NOT NULL" in sql
    assert "CONSTRAINT fk_operation_import_batch" in sql
    assert "REFERENCES drrk.import_batch (id)" in sql
    assert "CONSTRAINT uq_operation_batch_row UNIQUE" in sql
    assert "INSERT INTO drrk.operation_type" in sql