FinDataHub/Scripts/bin/tests/test_check_db.py

# coding=utf-8
"""check_db.py 数据库状态检测单元测试.

重点覆盖：
- tushare_data 数据库存在性检测（b501e35 新增逻辑）
- 三库（news_flash / akshare_data / tushare_data）表枚举
- ready 标志判定逻辑
- MySQL 连接失败时的错误降级
"""

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 check_db


class _FakeCursor:
    """模拟 pymysql cursor，按 SQL 关键词返回预设结果."""

    def __init__(self, databases=None, tables_map=None):
        self._databases = databases or []
        self._tables_map = tables_map or {}
        self._current_db = None

    def execute(self, sql, params=None):
        sql_upper = sql.upper().strip()
        if sql_upper.startswith("SHOW DATABASES"):
            self._fetchall_result = [(d,) for d in self._databases]
        elif sql_upper.startswith("SHOW TABLES"):
            self._fetchall_result = [
                (t,) for t in self._tables_map.get(self._current_db, [])
            ]

    def fetchall(self):
        return getattr(self, "_fetchall_result", [])

    def close(self):
        pass


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

    def __init__(self, databases=None, tables_map=None):
        self._cursor = _FakeCursor(databases, tables_map)
        self.selected_dbs = []

    def cursor(self):
        return self._cursor

    def select_db(self, db):
        self.selected_dbs.append(db)
        self._cursor._current_db = db

    def close(self):
        pass


class TestCheckDbTushare:
    """tushare_data 数据库检测测试（b501e35 新增）."""

    def test_tushare_exists_with_tables(self, capsys):
        """tushare_data 库存在且有表时，tushare_data_exists=True 且表被枚举."""
        fake_conn = _FakeConn(
            databases=["news_flash", "akshare_data", "tushare_data"],
            tables_map={
                "news_flash": ["news_sources", "news_flash", "push_log", "fetch_log"],
                "akshare_data": ["stock_daily"],
                "tushare_data": ["stock_daily", "stock_weekly"],
            },
        )
        with patch.object(check_db, "connect", return_value=fake_conn):
            check_db.main()

        result = json.loads(capsys.readouterr().out)
        assert result["tushare_data_exists"] is True
        assert "stock_daily" in result["tushare_data_tables"]
        assert "stock_weekly" in result["tushare_data_tables"]

    def test_tushare_not_exists(self, capsys):
        """tushare_data 库不存在时，tushare_data_exists=False 且表列表为空."""
        fake_conn = _FakeConn(
            databases=["news_flash", "akshare_data"],
            tables_map={
                "news_flash": ["news_sources", "news_flash", "push_log", "fetch_log"],
                "akshare_data": ["stock_daily"],
            },
        )
        with patch.object(check_db, "connect", return_value=fake_conn):
            check_db.main()

        result = json.loads(capsys.readouterr().out)
        assert result["tushare_data_exists"] is False
        assert result["tushare_data_tables"] == []

    def test_all_three_databases_checked(self, capsys):
        """三库同时存在时均被检查且 select_db 被调用三次."""
        fake_conn = _FakeConn(
            databases=["news_flash", "akshare_data", "tushare_data"],
            tables_map={
                "news_flash": ["news_sources", "news_flash", "push_log", "fetch_log"],
                "akshare_data": ["stock_daily"],
                "tushare_data": ["_table_meta"],
            },
        )
        with patch.object(check_db, "connect", return_value=fake_conn):
            check_db.main()

        assert "news_flash" in fake_conn.selected_dbs
        assert "akshare_data" in fake_conn.selected_dbs
        assert "tushare_data" in fake_conn.selected_dbs


class TestCheckDbReady:
    """ready 标志判定逻辑测试."""

    def test_ready_when_news_flash_has_all_required_tables(self, capsys):
        """news_flash 含全部必需表时 ready=True."""
        fake_conn = _FakeConn(
            databases=["news_flash", "akshare_data", "tushare_data"],
            tables_map={
                "news_flash": ["news_sources", "news_flash", "push_log", "fetch_log"],
                "akshare_data": [],
                "tushare_data": [],
            },
        )
        with patch.object(check_db, "connect", return_value=fake_conn):
            check_db.main()

        result = json.loads(capsys.readouterr().out)
        assert result["ready"] is True

    def test_not_ready_when_missing_required_table(self, capsys):
        """news_flash 缺少必需表时 ready=False."""
        fake_conn = _FakeConn(
            databases=["news_flash", "akshare_data", "tushare_data"],
            tables_map={
                "news_flash": ["news_sources", "news_flash"],  # 缺 push_log, fetch_log
                "akshare_data": [],
                "tushare_data": [],
            },
        )
        with patch.object(check_db, "connect", return_value=fake_conn):
            check_db.main()

        result = json.loads(capsys.readouterr().out)
        assert result["ready"] is False


class TestCheckDbConnectionError:
    """连接失败时的错误降级测试."""

    def test_connection_error_sets_error(self, capsys):
        """MySQL 不可达时 error 被设置，ready=False."""
        with patch.object(check_db, "connect", side_effect=RuntimeError("conn refused")):
            check_db.main()

        result = json.loads(capsys.readouterr().out)
        assert result["mysql_reachable"] is False
        assert result["ready"] is False
        assert "conn refused" in result["error"]

    def test_result_has_all_keys_even_on_error(self, capsys):
        """连接失败时结果仍包含所有预期字段（含 tushare 字段）."""
        with patch.object(check_db, "connect", side_effect=RuntimeError("timeout")):
            check_db.main()

        result = json.loads(capsys.readouterr().out)
        expected_keys = {
            "mysql_reachable", "news_flash_exists", "news_flash_tables",
            "akshare_data_exists", "akshare_data_tables",
            "tushare_data_exists", "tushare_data_tables",
            "ready", "error",
        }
        assert expected_keys.issubset(set(result.keys()))