FinDataHub/Scripts/bin/tests/test_backfill.py

# coding=utf-8
"""backfill.py 连接归还回归测试.

缺陷（2026-09-19 修复）
----------------------
`backfill.py` 的三处 `get_conn()` 都把 `put_conn(conn)` 放在**成功路径**：

    conn = get_conn()
    cur = conn.cursor()
    cur.execute(...)          # ← 这里抛错
    put_conn(conn)            # ← 永远执行不到

连接池有容量上限，异常路径不归还就会**永久占住池位**；反复失败即耗尽池，
表现为后续所有取连接的操作全部阻塞（详见 db_pool 的取连接超时修复）。

本测试锁住契约：**无论正常还是异常，put_conn 都必须被调用恰好一次**。
"""
import os
import sys
from pathlib import Path
from unittest.mock import MagicMock

import pytest

os.environ.setdefault("MYSQL_PASSWORD", "")

_bin_dir = os.path.join(os.path.dirname(__file__), "..")
if _bin_dir not in sys.path:
    sys.path.insert(0, _bin_dir)

# ── 与 test_newsflash_collect.py 同款：导入前预注入 mock，避免顶层依赖失败 ──
import importlib.util

_orig_sys_modules = dict(sys.modules)

_mock_common = MagicMock()
_mock_common.KX_HOME = "/tmp"
_mock_common.KXS_HOME = "/tmp"
_mock_common.BUNDLED_KXS_HOME = None
_mock_common.BIN_DIR = "/tmp"
_mock_common.MYSQL_CFG = {}
_mock_common.database_url = "mysql://localhost"
_mock_common.setup_kuaixun_paths = MagicMock()
_mock_common.load_env = MagicMock()

_mock_db_pool = MagicMock()
_mock_extend = MagicMock()
_mock_extend_fetchers = MagicMock()
_mock_extend_fetchers_manager = MagicMock()

_patches = {
    "common": _mock_common,
    "pymysql": MagicMock(),
    "db_pool": _mock_db_pool,
    "extend": _mock_extend,
    "extend.fetchers": _mock_extend_fetchers,
    "extend.fetchers.manager": _mock_extend_fetchers_manager,
}

try:
    sys.modules.update(_patches)
    _spec = importlib.util.spec_from_file_location(
        "backfill_mod", os.path.join(_bin_dir, "backfill.py")
    )
    _bf = importlib.util.module_from_spec(_spec)
    sys.modules["backfill_mod"] = _bf
    _spec.loader.exec_module(_bf)
finally:
    for _k, _v in _orig_sys_modules.items():
        sys.modules[_k] = _v
    for _k in _patches:
        if _k not in _orig_sys_modules:
            sys.modules.pop(_k, None)

backfill = _bf


class _FakeCursor:
    def __init__(self, rows=None, raise_on_execute=False):
        self._rows = rows or []
        self._raise = raise_on_execute

    def execute(self, sql, params=None):
        if self._raise:
            raise RuntimeError("模拟数据库错误")

    def fetchall(self):
        return list(self._rows)

    def close(self):
        pass

    @property
    def rowcount(self):
        return 0


class _FakeConn:
    def __init__(self, cur):
        self._cur = cur

    def cursor(self):
        return self._cur

    def commit(self):
        pass


@pytest.fixture
def pool(monkeypatch):
    """替换 get_conn/put_conn，记录调用次数。"""
    calls = {"get": 0, "put": 0}

    def _get():
        calls["get"] += 1
        return _FakeConn(_FakeCursor())

    def _put(conn):
        calls["put"] += 1

    monkeypatch.setattr(backfill, "get_conn", _get)
    monkeypatch.setattr(backfill, "put_conn", _put)
    return calls


class TestConnectionReturned:
    def test_get_source_map_returns_conn_on_success(self, pool, monkeypatch):
        conn = _FakeConn(_FakeCursor(rows=[(1, "cls")]))
        monkeypatch.setattr(backfill, "get_conn", lambda: conn)
        backfill.get_source_map()
        assert pool["put"] == 1

    def test_get_source_map_returns_conn_on_error(self, pool, monkeypatch):
        """核心回归：execute 抛错时连接仍须归还."""
        conn = _FakeConn(_FakeCursor(raise_on_execute=True))
        monkeypatch.setattr(backfill, "get_conn", lambda: conn)
        with pytest.raises(RuntimeError):
            backfill.get_source_map()
        assert pool["put"] == 1, "异常路径未归还连接 → 池位泄漏"

    def test_get_existing_ids_returns_conn_on_error(self, pool, monkeypatch):
        conn = _FakeConn(_FakeCursor(raise_on_execute=True))
        monkeypatch.setattr(backfill, "get_conn", lambda: conn)
        backfill.get_existing_ids("cls", 0, 1)      # 内部吞掉异常
        assert pool["put"] == 1, "异常路径未归还连接 → 池位泄漏"

    def test_get_existing_ids_returns_conn_on_success(self, pool, monkeypatch):
        conn = _FakeConn(_FakeCursor(rows=[("n1",)]))
        monkeypatch.setattr(backfill, "get_conn", lambda: conn)
        backfill.get_existing_ids("cls", 0, 1)
        assert pool["put"] == 1

    def test_insert_news_returns_conn_on_error(self, pool, monkeypatch):
        """cur.close() 抛错时也必须归还（finally 覆盖）. """
        cur = _FakeCursor()
        cur.close = MagicMock(side_effect=RuntimeError("close 失败"))
        conn = _FakeConn(cur)
        monkeypatch.setattr(backfill, "get_conn", lambda: conn)
        with pytest.raises(RuntimeError):
            backfill.insert_news([])
        assert pool["put"] == 1

    def test_insert_news_returns_conn_on_success(self, pool, monkeypatch):
        conn = _FakeConn(_FakeCursor())
        monkeypatch.setattr(backfill, "get_conn", lambda: conn)
        backfill.insert_news([])
        assert pool["put"] == 1