FinDataHub/Scripts/bin/tests/test_init_db.py

# coding=utf-8
"""init_db.py 数据库初始化脚本单元测试.

重点覆盖：
- init_tushare_data：tushare_data 库创建 + _table_meta 建表（b501e35 新增）
- init_tushare_data 被 main() 正确调用并纳入 steps
- execute_sql：成功/异常路径
- check_database：连接成功/失败
"""

import json
import os
import sys
from unittest.mock import MagicMock, patch

import pytest

# 将 Scripts/bin 加入 sys.path
_bin_dir = os.path.join(os.path.dirname(__file__), "..")
if _bin_dir not in sys.path:
    sys.path.insert(0, _bin_dir)

import init_db


class _FakeCursor:
    """模拟 cursor，记录执行的 SQL."""

    def __init__(self):
        self.executed = []

    def execute(self, sql, params=None):
        self.executed.append(sql)

    def close(self):
        pass


class _FakeConn:
    """模拟 pymysql 连接，记录 select_db 调用."""

    def __init__(self):
        self._cursor = _FakeCursor()
        self.selected_dbs = []
        self.committed = False

    def cursor(self):
        # execute_sql 用 with conn.cursor() as cur 模式
        return _ContextCursor(self._cursor)

    def select_db(self, db):
        self.selected_dbs.append(db)

    def commit(self):
        self.committed = True

    def close(self):
        pass


class _ContextCursor:
    """with conn.cursor() as cur 的上下文管理器包装."""

    def __init__(self, cursor):
        self._cursor = cursor

    def __enter__(self):
        return self._cursor

    def __exit__(self, *a):
        return False


# ── init_tushare_data ─────────────────────────────────────


class TestInitTushareData:
    """init_tushare_data 测试（b501e35 新增）."""

    def test_creates_database_and_table(self):
        """验证创建 tushare_data 库和 _table_meta 表."""
        conn = _FakeConn()
        steps = init_db.init_tushare_data(conn)

        assert len(steps) == 2
        assert "tushare_data" in steps[0]
        assert "_table_meta" in steps[1]

        # 验证 SQL 执行内容
        all_sql = " ".join(conn._cursor.executed)
        assert "CREATE DATABASE IF NOT EXISTS" in all_sql
        assert "tushare_data" in all_sql
        assert "CREATE TABLE IF NOT EXISTS" in all_sql
        assert "_table_meta" in all_sql

    def test_selects_tushare_database(self):
        """验证 select_db 被调用且参数为 tushare_data."""
        conn = _FakeConn()
        init_db.init_tushare_data(conn)

        assert init_db.MYSQL_DATABASE_TUSHARE in conn.selected_dbs

    def test_table_meta_has_correct_columns(self):
        """验证 _table_meta 表结构包含必要列."""
        conn = _FakeConn()
        init_db.init_tushare_data(conn)

        table_sql = conn._cursor.executed[1]
        for col in ["table_name", "last_update", "row_count", "status"]:
            assert col in table_sql

    def test_table_meta_primary_key(self):
        """验证 _table_meta 以 table_name 为主键."""
        conn = _FakeConn()
        init_db.init_tushare_data(conn)

        table_sql = conn._cursor.executed[1]
        assert "PRIMARY KEY" in table_sql
        assert "table_name" in table_sql


# ── execute_sql ───────────────────────────────────────────


class TestExecuteSql:
    """execute_sql 成功/异常路径测试."""

    def test_success_returns_true(self):
        conn = _FakeConn()
        assert init_db.execute_sql(conn, "CREATE DATABASE x") is True
        assert conn.committed is True

    def test_exception_returns_false(self, capsys):
        """执行失败时返回 False 并输出错误到 stderr."""
        class FailConn(_FakeConn):
            def cursor(self):
                raise RuntimeError("boom")

        conn = FailConn()
        result = init_db.execute_sql(conn, "SELECT 1")
        assert result is False
        captured = capsys.readouterr()
        assert "SQL ERROR" in captured.err


# ── check_database ────────────────────────────────────────


class TestCheckDatabase:
    """check_database 连接检查测试."""

    def test_success_returns_true(self):
        fake_conn = MagicMock()
        with patch.object(init_db, "connect", return_value=fake_conn):
            assert init_db.check_database("h", 3306, "u", "p", "db") is True
        fake_conn.close.assert_called_once()

    def test_failure_returns_false(self, capsys):
        with patch.object(init_db, "connect", side_effect=RuntimeError("refused")):
            assert init_db.check_database("h", 3306, "u", "p", "db") is False
        captured = capsys.readouterr()
        assert "CHECK FAILED" in captured.err


# ── main 集成 ─────────────────────────────────────────────


class TestMainIntegration:
    """main() 调用三个 init_* 函数的集成测试."""

    def test_main_includes_tushare_steps(self, capsys):
        """main() 结果中包含 tushare_data 初始化步骤."""
        fake_conn = _FakeConn()
        with patch.object(init_db, "connect", return_value=fake_conn):
            init_db.main()

        result = json.loads(capsys.readouterr().out)
        assert result["success"] is True
        all_steps = " ".join(result["steps"])
        assert "tushare_data" in all_steps
        assert "_table_meta" in all_steps

    def test_main_calls_all_three_init_functions(self, capsys):
        """main() 调用 news_flash + akshare_data + tushare_data 三个初始化."""
        fake_conn = _FakeConn()
        with patch.object(init_db, "connect", return_value=fake_conn):
            init_db.main()

        result = json.loads(capsys.readouterr().out)
        all_steps = " ".join(result["steps"])
        assert "news_flash" in all_steps
        assert "akshare_data" in all_steps
        assert "tushare_data" in all_steps

    def test_main_connection_error(self, capsys):
        """main() 连接失败时 success=False 且输出错误."""
        with patch.object(init_db, "connect", side_effect=RuntimeError("no mysql")):
            init_db.main()

        result = json.loads(capsys.readouterr().out)
        assert result["success"] is False
        assert "no mysql" in result["error"]