FinDataHub
11d36d1c1e
问题 - 索引中该目录记为小写 `scripts/`，工作区实为 `Scripts/`。因 core.ignorecase=true， 本地 `git diff -- Scripts/bin/x.py` 会**静默返回空**（看似无改动）， 且在大小写敏感环境（Linux CI / Docker 挂载）会检出 `scripts/`， 与全仓所有 `Scripts/...` 引用不匹配。 - 该目录此前 10 个脚本（check_python_version.py、perf_snapshot.py、 install_collector_launchd.py、mysql_toggle.sh 等）从未入库， 而 CI 第 31 行正调用 `Scripts/bin/check_python_version.py --strict` —— 全新检出必然失败。 改动 - 索引路径统一为 `Scripts/`（工作区文件未移动，仅索引改名）。 - 补齐上述 10 个未跟踪脚本及其测试。 影响 - 修正前索引 35 条（全为小写路径），修正后 45 条。
449 lines
14 KiB
Python
449 lines
14 KiB
Python
# coding=utf-8
|
||
"""db_pool.py 轻量 MySQL 连接池单元测试.
|
||
|
||
重点覆盖：
|
||
- _connect_cfg：超时参数注入
|
||
- _init_pool：连接池初始化
|
||
- get_conn：惰性创建 + 自动重连 + 重试机制
|
||
- put_conn：归还连接（池满时关闭连接）
|
||
- DBContext：上下文管理器（自动获取/归还）
|
||
"""
|
||
|
||
import pytest
|
||
from unittest.mock import patch, MagicMock, call
|
||
from queue import Queue
|
||
|
||
# 将 Scripts/bin 加入 sys.path
|
||
import os
|
||
import sys
|
||
|
||
_bin_dir = os.path.join(os.path.dirname(__file__), "..")
|
||
if _bin_dir not in sys.path:
|
||
sys.path.insert(0, _bin_dir)
|
||
|
||
|
||
# ── _connect_cfg ─────────────────────────────────────────
|
||
|
||
|
||
class TestConnectCfg:
|
||
"""连接配置构造测试."""
|
||
|
||
@patch("db_pool.MYSQL_CFG", {"host": "localhost", "port": 3306, "user": "root", "password": "", "database": "test", "charset": "utf8mb4"})
|
||
def test_returns_dict_with_timeouts(self):
|
||
"""返回的配置包含超时参数."""
|
||
from db_pool import _connect_cfg, _CONNECT_TIMEOUT, _READ_TIMEOUT
|
||
|
||
cfg = _connect_cfg()
|
||
assert cfg["connect_timeout"] == _CONNECT_TIMEOUT
|
||
assert cfg["read_timeout"] == _READ_TIMEOUT
|
||
|
||
@patch("db_pool.MYSQL_CFG", {"host": "localhost", "port": 3306, "user": "root", "password": "", "database": "test", "charset": "utf8mb4"})
|
||
def test_preserves_original_cfg_values(self):
|
||
"""保留原始 MYSQL_CFG 的值."""
|
||
from db_pool import _connect_cfg
|
||
|
||
cfg = _connect_cfg()
|
||
assert cfg["host"] == "localhost"
|
||
assert cfg["port"] == 3306
|
||
assert cfg["user"] == "root"
|
||
assert cfg["database"] == "test"
|
||
|
||
@patch("db_pool.MYSQL_CFG", {"host": "localhost", "port": 3306, "user": "root", "password": "", "database": "test", "charset": "utf8mb4", "connect_timeout": 99})
|
||
def test_setdefault_does_not_override(self):
|
||
"""setdefault 不覆盖已有 connect_timeout."""
|
||
from db_pool import _connect_cfg
|
||
|
||
cfg = _connect_cfg()
|
||
# 已有 connect_timeout=99 时 setdefault 不覆盖
|
||
assert cfg["connect_timeout"] == 99
|
||
|
||
|
||
# ── _init_pool ───────────────────────────────────────────
|
||
|
||
|
||
class TestInitPool:
|
||
"""连接池初始化测试."""
|
||
|
||
def test_init_creates_pool_with_none_slots(self):
|
||
"""初始化后池中填充 None（惰性创建）."""
|
||
import db_pool
|
||
|
||
db_pool._pool = None
|
||
db_pool._init_pool(max_conn=4)
|
||
assert db_pool._pool is not None
|
||
assert db_pool._pool.qsize() == 4
|
||
# 所有槽位应为 None
|
||
for _ in range(4):
|
||
assert db_pool._pool.get() is None
|
||
|
||
def test_init_idempotent(self):
|
||
"""重复调用 _init_pool 不重新创建池."""
|
||
import db_pool
|
||
|
||
db_pool._pool = None
|
||
db_pool._init_pool(max_conn=4)
|
||
first_pool = db_pool._pool
|
||
db_pool._init_pool(max_conn=8)
|
||
assert db_pool._pool is first_pool # 同一对象
|
||
|
||
def teardown_method(self, method):
|
||
"""每个测试后重置全局 _pool."""
|
||
import db_pool
|
||
|
||
db_pool._pool = None
|
||
|
||
|
||
# ── get_conn ─────────────────────────────────────────────
|
||
|
||
|
||
class TestGetConn:
|
||
"""连接获取测试 — 惰性创建 + 重连 + 重试."""
|
||
|
||
def teardown_method(self, method):
|
||
"""每个测试后重置全局 _pool."""
|
||
import db_pool
|
||
|
||
db_pool._pool = None
|
||
|
||
@patch("db_pool.pymysql.connect")
|
||
def test_creates_new_connection_when_none(self, mock_connect):
|
||
"""池中为 None 时创建新连接."""
|
||
import db_pool
|
||
|
||
mock_conn = MagicMock()
|
||
mock_connect.return_value = mock_conn
|
||
|
||
db_pool._pool = None
|
||
db_pool._init_pool(max_conn=2)
|
||
conn = db_pool.get_conn()
|
||
|
||
assert conn is mock_conn
|
||
mock_connect.assert_called_once()
|
||
|
||
@patch("db_pool.pymysql.connect")
|
||
def test_reuses_existing_connection_via_ping(self, mock_connect):
|
||
"""池中有连接时通过 ping(reconnect=True) 验证可用性."""
|
||
import db_pool
|
||
|
||
mock_conn = MagicMock()
|
||
mock_conn.open = True
|
||
|
||
db_pool._pool = None
|
||
db_pool._init_pool(max_conn=2)
|
||
# 清空池中所有 None，只放入 mock_conn
|
||
db_pool._pool.get()
|
||
db_pool._pool.get()
|
||
db_pool._pool.put(mock_conn)
|
||
|
||
conn = db_pool.get_conn()
|
||
|
||
assert conn is mock_conn
|
||
mock_conn.ping.assert_called_once_with(reconnect=True)
|
||
# 不应调用 pymysql.connect
|
||
mock_connect.assert_not_called()
|
||
|
||
@patch("db_pool.pymysql.connect")
|
||
def test_retries_on_failure(self, mock_connect):
|
||
"""连接失败时重试 max_retries 次."""
|
||
import db_pool
|
||
|
||
mock_connect.side_effect = Exception("connection refused")
|
||
|
||
db_pool._pool = None
|
||
db_pool._init_pool(max_conn=2)
|
||
|
||
with patch("time.sleep"): # 跳过重试延迟
|
||
with pytest.raises(Exception, match="connection refused"):
|
||
db_pool.get_conn(max_retries=3, retry_delay=0)
|
||
|
||
assert mock_connect.call_count == 3
|
||
|
||
@patch("db_pool.pymysql.connect")
|
||
def test_reconnect_after_ping_failure(self, mock_connect):
|
||
"""ping 失败后重新创建连接."""
|
||
import db_pool
|
||
|
||
mock_conn = MagicMock()
|
||
mock_conn.ping.side_effect = Exception("connection lost")
|
||
new_conn = MagicMock()
|
||
mock_connect.return_value = new_conn
|
||
|
||
db_pool._pool = None
|
||
db_pool._init_pool(max_conn=2)
|
||
# 第一次从池中拿到 mock_conn（ping 失败）
|
||
# 第二次从池中拿到 None（新建 new_conn）
|
||
db_pool._pool = Queue(maxsize=2)
|
||
db_pool._pool.put(mock_conn)
|
||
db_pool._pool.put(None)
|
||
|
||
with patch("time.sleep"):
|
||
conn = db_pool.get_conn(max_retries=3, retry_delay=0)
|
||
|
||
assert conn is new_conn
|
||
mock_conn.ping.assert_called_once()
|
||
|
||
@patch("db_pool.pymysql.connect")
|
||
def test_init_pool_called_if_none(self, mock_connect):
|
||
"""_pool 为 None 时自动初始化."""
|
||
import db_pool
|
||
|
||
mock_connect.return_value = MagicMock()
|
||
|
||
db_pool._pool = None
|
||
db_pool.get_conn()
|
||
|
||
assert db_pool._pool is not None
|
||
|
||
|
||
# ── put_conn ─────────────────────────────────────────────
|
||
|
||
|
||
class TestPutConn:
|
||
"""连接归还测试."""
|
||
|
||
def teardown_method(self, method):
|
||
"""每个测试后重置全局 _pool."""
|
||
import db_pool
|
||
|
||
db_pool._pool = None
|
||
|
||
def test_put_open_connection_back_to_pool(self):
|
||
"""开放连接归还到池中."""
|
||
import db_pool
|
||
|
||
mock_conn = MagicMock()
|
||
mock_conn.open = True
|
||
|
||
db_pool._pool = None
|
||
db_pool._init_pool(max_conn=4)
|
||
# 取出一个 None 腾出位置，避免池满导致连接被关闭
|
||
db_pool._pool.get()
|
||
db_pool.put_conn(mock_conn)
|
||
|
||
# 池应已满（3 个 None + 1 个 mock_conn）
|
||
assert db_pool._pool.qsize() == 4
|
||
# Queue 是 FIFO，取出的前 3 个是 None，第 4 个是 mock_conn
|
||
assert db_pool._pool.get() is None
|
||
assert db_pool._pool.get() is None
|
||
assert db_pool._pool.get() is None
|
||
assert db_pool._pool.get() is mock_conn
|
||
|
||
def test_put_closed_connection_as_none(self):
|
||
"""已关闭连接归还为 None（惰性标记）."""
|
||
import db_pool
|
||
|
||
mock_conn = MagicMock()
|
||
mock_conn.open = False
|
||
|
||
db_pool._pool = None
|
||
db_pool._init_pool(max_conn=4)
|
||
db_pool.put_conn(mock_conn)
|
||
|
||
assert db_pool._pool.get() is None
|
||
|
||
def test_put_none_connection(self):
|
||
"""None 连接直接归还."""
|
||
import db_pool
|
||
|
||
db_pool._pool = None
|
||
db_pool._init_pool(max_conn=4)
|
||
db_pool.put_conn(None)
|
||
|
||
assert db_pool._pool.get() is None
|
||
|
||
def test_pool_full_closes_connection(self):
|
||
"""池满时关闭连接而非阻塞."""
|
||
import db_pool
|
||
|
||
mock_conn = MagicMock()
|
||
mock_conn.open = True
|
||
|
||
db_pool._pool = None
|
||
db_pool._init_pool(max_conn=2)
|
||
# 池已满（2个 None）
|
||
# put_conn 时池满，应关闭连接
|
||
extra_conn = MagicMock()
|
||
extra_conn.open = True
|
||
db_pool.put_conn(extra_conn)
|
||
|
||
# 连接被关闭
|
||
extra_conn.close.assert_called_once()
|
||
|
||
def test_put_conn_initializes_pool_if_none(self):
|
||
"""_pool 为 None 时自动初始化."""
|
||
import db_pool
|
||
|
||
db_pool._pool = None
|
||
db_pool.put_conn(None)
|
||
assert db_pool._pool is not None
|
||
|
||
|
||
# ── 回归：连接泄漏导致后端挂死 ───────────────────────────
|
||
#
|
||
# 曾存在的缺陷：_query_all 的异常路径不归还连接，而 _pool.get() 没有 timeout。
|
||
# 池容量 8，漏满之后所有请求永久阻塞在 get() 上——后端不报错、只剩线程堆积。
|
||
# 下面四条分别锁住：取连接有超时、归还有回滚、失败不吞池位、异常出栈回滚。
|
||
|
||
|
||
class TestPoolExhaustionRegression:
|
||
"""连接池耗尽与脏事务回归测试."""
|
||
|
||
def teardown_method(self, method):
|
||
import db_pool
|
||
|
||
db_pool._pool = None
|
||
|
||
def test_get_conn_times_out_instead_of_hanging(self, monkeypatch):
|
||
"""池被耗尽时必须抛 TimeoutError，而不是永久阻塞."""
|
||
import db_pool
|
||
|
||
# 容量 1 且槽位已被取走 → 池空
|
||
db_pool._pool = Queue(maxsize=1)
|
||
monkeypatch.setattr(db_pool, "_POOL_GET_TIMEOUT", 0.05)
|
||
|
||
with pytest.raises(TimeoutError, match="连接池已耗尽"):
|
||
db_pool.get_conn()
|
||
|
||
def test_put_conn_rolls_back_dirty_transaction(self):
|
||
"""归还未提交事务的连接前必须 rollback."""
|
||
import db_pool
|
||
|
||
mock_conn = MagicMock()
|
||
mock_conn.open = True
|
||
|
||
db_pool._pool = None
|
||
db_pool._init_pool(max_conn=4)
|
||
db_pool._pool.get() # 腾出一格
|
||
db_pool.put_conn(mock_conn)
|
||
|
||
mock_conn.rollback.assert_called_once()
|
||
assert db_pool._pool.get() is None
|
||
# mock_conn 已归还（3 个 None 之后是它）
|
||
assert db_pool._pool.get() is None
|
||
assert db_pool._pool.get() is None
|
||
assert db_pool._pool.get() is mock_conn
|
||
|
||
def test_put_conn_discards_connection_when_rollback_fails(self):
|
||
"""rollback 失败说明连接已坏，应丢弃（放 None）而非塞回池."""
|
||
import db_pool
|
||
|
||
mock_conn = MagicMock()
|
||
mock_conn.open = True
|
||
mock_conn.rollback.side_effect = Exception("connection gone")
|
||
|
||
db_pool._pool = None
|
||
db_pool._init_pool(max_conn=2)
|
||
db_pool._pool.get() # 腾出一格
|
||
db_pool.put_conn(mock_conn)
|
||
|
||
mock_conn.close.assert_called_once()
|
||
# 池里补回的是 None，不是坏连接
|
||
assert db_pool._pool.get() is None
|
||
|
||
@patch("db_pool.pymysql.connect")
|
||
def test_get_conn_returns_slot_on_total_failure(self, mock_connect):
|
||
"""连接创建全部失败后，占用的池位必须归还，否则池会被逐步耗尽."""
|
||
import db_pool
|
||
|
||
mock_connect.side_effect = Exception("connection refused")
|
||
|
||
db_pool._pool = None
|
||
db_pool._init_pool(max_conn=2)
|
||
|
||
with patch("time.sleep"):
|
||
with pytest.raises(Exception, match="connection refused"):
|
||
db_pool.get_conn(max_retries=2, retry_delay=0)
|
||
|
||
# 池容量未被吃掉：取走 2 个仍能拿到（不会 Empty 阻塞）
|
||
assert db_pool._pool.get(timeout=1) is None
|
||
assert db_pool._pool.get(timeout=1) is None
|
||
|
||
@patch("db_pool.get_conn")
|
||
@patch("db_pool.put_conn")
|
||
def test_dbcontext_rolls_back_on_exception(self, mock_put, mock_get):
|
||
"""异常出栈时 DBContext 显式回滚，避免脏事务被复用."""
|
||
import db_pool
|
||
|
||
mock_conn = MagicMock()
|
||
mock_get.return_value = mock_conn
|
||
|
||
with pytest.raises(ValueError):
|
||
with db_pool.DBContext():
|
||
raise ValueError("boom")
|
||
|
||
mock_conn.rollback.assert_called_once()
|
||
mock_put.assert_called_once_with(mock_conn)
|
||
|
||
@patch("db_pool.get_conn")
|
||
@patch("db_pool.put_conn")
|
||
def test_dbcontext_no_rollback_on_success(self, mock_put, mock_get):
|
||
"""正常退出时不额外回滚（由 put_conn 统一兜底）."""
|
||
import db_pool
|
||
|
||
mock_conn = MagicMock()
|
||
mock_get.return_value = mock_conn
|
||
|
||
with db_pool.DBContext():
|
||
pass
|
||
|
||
mock_conn.rollback.assert_not_called()
|
||
mock_put.assert_called_once_with(mock_conn)
|
||
|
||
|
||
# ── DBContext ────────────────────────────────────────────
|
||
|
||
|
||
class TestDBContext:
|
||
"""上下文管理器测试."""
|
||
|
||
def teardown_method(self, method):
|
||
"""每个测试后重置全局 _pool."""
|
||
import db_pool
|
||
|
||
db_pool._pool = None
|
||
|
||
@patch("db_pool.get_conn")
|
||
@patch("db_pool.put_conn")
|
||
def test_context_manager_get_and_put(self, mock_put, mock_get):
|
||
"""DBContext 正确获取和归还连接."""
|
||
import db_pool
|
||
|
||
mock_conn = MagicMock()
|
||
mock_get.return_value = mock_conn
|
||
|
||
with db_pool.DBContext() as conn:
|
||
assert conn is mock_conn
|
||
|
||
mock_get.assert_called_once()
|
||
mock_put.assert_called_once_with(mock_conn)
|
||
|
||
@patch("db_pool.get_conn")
|
||
@patch("db_pool.put_conn")
|
||
def test_context_manager_returns_connection_even_on_exception(self, mock_put, mock_get):
|
||
"""异常时连接也被归还."""
|
||
import db_pool
|
||
|
||
mock_conn = MagicMock()
|
||
mock_get.return_value = mock_conn
|
||
|
||
with pytest.raises(ValueError):
|
||
with db_pool.DBContext() as conn:
|
||
raise ValueError("test error")
|
||
|
||
mock_put.assert_called_once_with(mock_conn)
|
||
|
||
@patch("db_pool.get_conn")
|
||
@patch("db_pool.put_conn")
|
||
def test_context_manager_sets_conn_none_after_exit(self, mock_put, mock_get):
|
||
"""退出后 self.conn 被置为 None."""
|
||
import db_pool
|
||
|
||
mock_conn = MagicMock()
|
||
mock_get.return_value = mock_conn
|
||
|
||
ctx = db_pool.DBContext()
|
||
with ctx as conn:
|
||
pass
|
||
|
||
assert ctx.conn is None
|