claude_dev
c041227139
data/ 规则匹配任意 data 目录，误伤 tests/data（硬化测试）和 docs/data（15min设计文档）， 此前靠 git add -f 临时绕。加 !tests/data/ !docs/data/ 例外一劳永逸。
235 lines
8.6 KiB
Python
235 lines
8.6 KiB
Python
"""
|
||
硬化测试：backfill_15min_baostock.py 三个硬伤的回归覆盖。
|
||
|
||
覆盖：
|
||
1. 异常时强制重登（_relogin 在 backfill_one 重试 loop 中被调用）
|
||
2. 断路器（连续 N 失败 → break + sys.exit(2)）
|
||
3. marker-based resume（load_progress 从 marker 文件构造 done_set）
|
||
4. _login_with_timeout SIGALRM 超时（防 bs.login 永久挂起）
|
||
|
||
mock baostock，不发起真实连接。
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import tempfile
|
||
import time
|
||
from pathlib import Path
|
||
|
||
import pandas as pd
|
||
import pytest
|
||
|
||
# 在 import backfill 模块前重定向 STOCK_ROOT，避免 setup_logging() 写入 NAS 挂载点
|
||
_ROOT = Path(__file__).resolve().parents[2]
|
||
if str(_ROOT) not in sys.path:
|
||
sys.path.insert(0, str(_ROOT))
|
||
_FAKE_ROOT = tempfile.mkdtemp(prefix="backfill_test_")
|
||
os.environ["STOCK_ROOT"] = _FAKE_ROOT
|
||
|
||
import scripts.data_platform.backfill_15min_baostock as bf # noqa: E402
|
||
|
||
|
||
# ======================== fixture ========================
|
||
|
||
@pytest.fixture
|
||
def fake_df():
|
||
"""一份合法的 15min DataFrame，够 backfill_one 写盘"""
|
||
return pd.DataFrame({
|
||
"day": ["2024-01-02 09:45:00", "2024-01-02 10:00:00"],
|
||
"open": [10.0, 10.2], "high": [10.3, 10.4],
|
||
"low": [9.9, 10.1], "close": [10.2, 10.3],
|
||
"volume": ["1000", "1200"], "amount": ["10000", "12000"],
|
||
})
|
||
|
||
|
||
# ======================== 1. 异常时强制重登 ========================
|
||
|
||
def test_relogin_called_on_exception(tmp_path, monkeypatch, fake_df):
|
||
"""
|
||
fetch_bs_15min 第一次抛 socket 异常 →
|
||
backfill_one 重试 loop 应调 _relogin() 强制重登（新 socket）→
|
||
第二次 fetch 成功 → status=ok。
|
||
"""
|
||
calls = {"fetch": 0, "relogin": 0}
|
||
|
||
def mock_fetch(bs_code, start_date, end_date):
|
||
calls["fetch"] += 1
|
||
if calls["fetch"] == 1:
|
||
raise OSError(57, "Socket is not connected") # 模拟 baostock 断连
|
||
return fake_df
|
||
|
||
def mock_relogin():
|
||
calls["relogin"] += 1
|
||
return True
|
||
|
||
monkeypatch.setattr(bf, "MINUTE_15_DIR", tmp_path)
|
||
monkeypatch.setattr(bf, "fetch_bs_15min", mock_fetch)
|
||
monkeypatch.setattr(bf, "_relogin", mock_relogin)
|
||
monkeypatch.setattr(bf, "is_backfilled", lambda p: False)
|
||
|
||
status, rows = bf.backfill_one("000001", "2024-01-01", "2026-01-01", force=True)
|
||
|
||
assert status == "ok"
|
||
assert rows == 2
|
||
assert calls["fetch"] == 2 # 重试了一次
|
||
assert calls["relogin"] == 1 # 异常分支强制重登一次
|
||
|
||
|
||
def test_relogin_failure_aborts_retry(tmp_path, monkeypatch):
|
||
"""
|
||
异常 + _relogin 返回 False → 直接放弃，不再空转重试。
|
||
"""
|
||
fetch_calls = {"n": 0}
|
||
|
||
def mock_fetch(*args, **kwargs):
|
||
fetch_calls["n"] += 1
|
||
raise OSError(57, "Socket is not connected")
|
||
|
||
monkeypatch.setattr(bf, "MINUTE_15_DIR", tmp_path)
|
||
monkeypatch.setattr(bf, "fetch_bs_15min", mock_fetch)
|
||
monkeypatch.setattr(bf, "_relogin", lambda: False)
|
||
monkeypatch.setattr(bf, "is_backfilled", lambda p: False)
|
||
|
||
status, rows = bf.backfill_one("000001", "2024-01-01", "2026-01-01", force=True)
|
||
|
||
assert status == "failed"
|
||
# _relogin 失败立即 break：fetch 只被调用 1 次（不会重试到 BS_MAX_RETRIES）
|
||
assert fetch_calls["n"] == 1
|
||
|
||
|
||
# ======================== 2. 断路器 ========================
|
||
|
||
def test_circuit_breaker_triggers(monkeypatch, tmp_path):
|
||
"""
|
||
主循环连续 N 只 failed → 断路器触发，sys.exit(2)。
|
||
不真实登录 baostock、不读 CSV。
|
||
"""
|
||
# 缩小阈值，加快测试
|
||
monkeypatch.setattr(bf, "CIRCUIT_BREAKER_THRESHOLD", 3)
|
||
# 所有桩
|
||
monkeypatch.setattr(bf, "NAS_ROOT", tmp_path) # NAS_ROOT.exists() 通过
|
||
monkeypatch.setattr(bf, "MINUTE_15_DIR", tmp_path / "15min")
|
||
monkeypatch.setattr(bf, "_login_with_timeout", lambda timeout=30: True)
|
||
monkeypatch.setattr(bf, "_relogin", lambda: True)
|
||
monkeypatch.setattr(bf, "bs", type("S", (), {
|
||
"logout": staticmethod(lambda: None),
|
||
"login": staticmethod(lambda: type("L", (), {"error_code": "0", "error_msg": ""})()),
|
||
}))
|
||
monkeypatch.setattr(bf, "backfill_one", lambda *a, **k: ("failed", 0))
|
||
monkeypatch.setattr(bf, "save_progress", lambda s: None)
|
||
monkeypatch.setattr(bf, "time", type("T", (), {
|
||
"sleep": staticmethod(lambda s: None),
|
||
"time": staticmethod(time.time),
|
||
}))
|
||
# argv：5 个 fake 代码，触发阈值 3 后应中途 break
|
||
monkeypatch.setattr(sys, "argv", [
|
||
"backfill_15min_baostock.py",
|
||
"--codes", "000001,000002,000003,000004,000005",
|
||
])
|
||
|
||
with pytest.raises(SystemExit) as exc_info:
|
||
bf.main()
|
||
|
||
assert exc_info.value.code == 2, "断路器触发应以退出码 2 退出"
|
||
|
||
|
||
def test_circuit_breaker_resets_on_ok(monkeypatch, tmp_path):
|
||
"""
|
||
ok 重置连续失败计数：N-1 失败后 1 个 ok，再 N-1 失败 → 不应触发断路。
|
||
最终正常退出（无 SystemExit）。
|
||
"""
|
||
monkeypatch.setattr(bf, "CIRCUIT_BREAKER_THRESHOLD", 3)
|
||
monkeypatch.setattr(bf, "NAS_ROOT", tmp_path)
|
||
monkeypatch.setattr(bf, "MINUTE_15_DIR", tmp_path / "15min")
|
||
monkeypatch.setattr(bf, "_login_with_timeout", lambda timeout=30: True)
|
||
monkeypatch.setattr(bf, "_relogin", lambda: True)
|
||
monkeypatch.setattr(bf, "bs", type("S", (), {
|
||
"logout": staticmethod(lambda: None),
|
||
"login": staticmethod(lambda: type("L", (), {"error_code": "0", "error_msg": ""})()),
|
||
}))
|
||
# 序列：fail, fail, ok, fail, fail → 连续计数峰 2 < 3，不触发
|
||
seq = [("failed", 0), ("failed", 0), ("ok", 1), ("failed", 0), ("failed", 0)]
|
||
monkeypatch.setattr(bf, "backfill_one", lambda *a, **k: seq.pop(0))
|
||
monkeypatch.setattr(bf, "save_progress", lambda s: None)
|
||
monkeypatch.setattr(bf, "time", type("T", (), {
|
||
"sleep": staticmethod(lambda s: None),
|
||
"time": staticmethod(time.time),
|
||
}))
|
||
monkeypatch.setattr(sys, "argv", [
|
||
"backfill_15min_baostock.py",
|
||
"--codes", "000001,000002,000003,000004,000005",
|
||
])
|
||
|
||
# 正常跑完不应抛 SystemExit
|
||
bf.main()
|
||
|
||
|
||
# ======================== 3. marker-based resume ========================
|
||
|
||
def test_load_progress_from_markers(tmp_path, monkeypatch):
|
||
"""
|
||
load_progress 从 .{prefix}{code}_15min.baostock marker 构造 done_set。
|
||
failed 票无 marker → 不在 done_set → 下次复跑会被重试。
|
||
"""
|
||
monkeypatch.setattr(bf, "MINUTE_15_DIR", tmp_path)
|
||
|
||
# 成功票：写 marker
|
||
for name in [".sh600000_15min.baostock", ".sz000001_15min.baostock",
|
||
".sh688981_15min.baostock"]:
|
||
(tmp_path / name).write_text("2026-07-13T00:00:00")
|
||
# 失败票：只有 parquet，无 marker（不应进 done_set）
|
||
(tmp_path / "sz300750_15min.parquet").write_text("dummy")
|
||
# 无关文件
|
||
(tmp_path / ". unrelated.baostock").write_text("x")
|
||
(tmp_path / ".sh600000_daily.baostock").write_text("x") # 非 _15min
|
||
|
||
done = bf.load_progress()
|
||
|
||
assert done == {"600000", "000001", "688981"}
|
||
assert "300750" not in done # failed（无 marker）会被重试
|
||
|
||
|
||
def test_load_progress_empty_when_dir_missing(tmp_path, monkeypatch):
|
||
"""目录不存在 → 空集合（不抛异常）"""
|
||
monkeypatch.setattr(bf, "MINUTE_15_DIR", tmp_path / "nonexistent")
|
||
assert bf.load_progress() == set()
|
||
|
||
|
||
# ======================== 4. login 超时 ========================
|
||
|
||
def test_login_with_timeout_fires(monkeypatch):
|
||
"""
|
||
bs.login() 卡死（sleep 远大于 timeout）→ SIGALRM 触发 → 返回 False。
|
||
验证不会永久挂起。
|
||
"""
|
||
def slow_login():
|
||
time.sleep(10) # 模拟 baostock 卡死
|
||
return type("L", (), {"error_code": "0", "error_msg": ""})()
|
||
|
||
monkeypatch.setattr(bf, "bs", type("S", (), {"login": staticmethod(slow_login)}))
|
||
|
||
t0 = time.time()
|
||
ok = bf._login_with_timeout(timeout=1)
|
||
elapsed = time.time() - t0
|
||
|
||
assert ok is False
|
||
assert elapsed < 3, f"超时应快速返回，实际 {elapsed:.1f}s"
|
||
|
||
|
||
def test_login_with_timeout_success(monkeypatch):
|
||
"""正常登录 → 返回 True"""
|
||
def quick_login():
|
||
return type("L", (), {"error_code": "0", "error_msg": ""})()
|
||
|
||
monkeypatch.setattr(bf, "bs", type("S", (), {"login": staticmethod(quick_login)}))
|
||
assert bf._login_with_timeout(timeout=5) is True
|
||
|
||
|
||
def test_login_with_timeout_login_failure(monkeypatch):
|
||
"""baostock 返回 error_code != 0 → 返回 False"""
|
||
def fail_login():
|
||
return type("L", (), {"error_code": "10001", "error_msg": "网络断开"})()
|
||
|
||
monkeypatch.setattr(bf, "bs", type("S", (), {"login": staticmethod(fail_login)}))
|
||
assert bf._login_with_timeout(timeout=5) is False
|