buffer_platform/tests/test_migration.py

"""
数据缓冲平台 - SQLite 到 MySQL 数据迁移脚本测试
"""
import logging
from datetime import datetime
from unittest.mock import MagicMock, patch

import pytest
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from app.models import Base, MarketData, SymbolTimestamp


def _create_sqlite_source(tmp_path):
    """创建包含测试数据的文件 SQLite 源数据库。"""
    db_path = tmp_path / "test_source.db"
    engine = create_engine(f"sqlite:///{db_path}")
    MarketData.__table__.create(bind=engine)
    SymbolTimestamp.__table__.create(bind=engine)
    Session = sessionmaker(bind=engine)
    session = Session()
    fetched_at = datetime(2026, 7, 4, 10, 0, 0)
    session.add_all(
        [
            MarketData(
                symbol="AG2606",
                data_type="futures",
                period="5min",
                candles_json='[{"open": 100, "close": 110}]',
                current_price=123.45,
                fetched_at=fetched_at,
                candle_count=1,
            ),
            MarketData(
                symbol="AG2606",
                data_type="futures",
                period="15min",
                candles_json='[{"open": 105, "close": 115}]',
                current_price=123.45,
                fetched_at=fetched_at,
                candle_count=1,
            ),
            SymbolTimestamp(
                symbol="AG2606",
                data_type="futures",
                last_refresh_at=fetched_at,
                refresh_count=5,
            ),
        ]
    )
    session.commit()
    session.close()
    return engine


def _make_mysql_session(market_data_count=0, symbol_timestamp_count=0):
    """构造模拟 MySQL 会话，支持记录写入和计数查询。"""
    mysql_session = MagicMock()
    mysql_session.__enter__ = MagicMock(return_value=mysql_session)
    mysql_session.__exit__ = MagicMock(return_value=False)

    added_market_data = []
    added_symbol_timestamps = []

    def mock_query(model):
        query = MagicMock()
        if model is MarketData:
            query.count.return_value = market_data_count
        elif model is SymbolTimestamp:
            query.count.return_value = symbol_timestamp_count
        return query

    mysql_session.query.side_effect = mock_query
    mysql_session.add_all.side_effect = lambda records: (
        added_market_data.extend(records)
        if records and isinstance(records[0], MarketData)
        else added_symbol_timestamps.extend(records)
    )

    return mysql_session, added_market_data, added_symbol_timestamps


class TestMigrateSqliteToMysql:
    """migrate_sqlite_to_mysql 行为测试。"""

    def test_migrates_data_when_mysql_tables_are_empty(self, caplog, tmp_path):
        """MySQL 表为空时，应从 SQLite 迁移全部数据。"""
        sqlite_engine = _create_sqlite_source(tmp_path)
        mysql_session, added_market_data, added_symbol_timestamps = _make_mysql_session(
            market_data_count=0, symbol_timestamp_count=0
        )

        with patch("app.migration.create_engine", return_value=sqlite_engine):
            with patch(
                "app.migration.MySQLSessionLocal", return_value=mysql_session
            ):
                with caplog.at_level(logging.INFO, logger="app.migration"):
                    from app.migration import migrate_sqlite_to_mysql

                    result = migrate_sqlite_to_mysql()

        assert result is True
        assert len(added_market_data) == 2
        assert len(added_symbol_timestamps) == 1
        assert added_market_data[0].symbol == "AG2606"
        assert added_market_data[0].period == "5min"
        assert added_symbol_timestamps[0].symbol == "AG2606"
        assert "迁移完成" in caplog.text
        assert "market_data: 2" in caplog.text
        assert "symbol_timestamps: 1" in caplog.text
        mysql_session.commit.assert_called_once()

    def test_skips_migration_when_mysql_has_data(self, caplog, tmp_path):
        """MySQL 表非空时，应跳过迁移并输出提示日志。"""
        sqlite_engine = _create_sqlite_source(tmp_path)
        mysql_session, added_market_data, added_symbol_timestamps = _make_mysql_session(
            market_data_count=1, symbol_timestamp_count=1
        )

        with patch("app.migration.create_engine", return_value=sqlite_engine):
            with patch(
                "app.migration.MySQLSessionLocal", return_value=mysql_session
            ):
                with caplog.at_level(logging.INFO, logger="app.migration"):
                    from app.migration import migrate_sqlite_to_mysql

                    result = migrate_sqlite_to_mysql()

        assert result is False
        assert len(added_market_data) == 0
        assert len(added_symbol_timestamps) == 0
        assert "MySQL 已存在数据" in caplog.text
        assert "跳过迁移" in caplog.text
        mysql_session.commit.assert_not_called()