SHA256
305 lines
13 KiB
Python
305 lines
13 KiB
Python
"""验证 daily.py 中 BaoStock 日线抓取逻辑（不打网络）。"""
|
|
|
|
from contextlib import contextmanager
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
import pandas as pd
|
|
|
|
from src.fetchers import daily
|
|
|
|
|
|
def _make_market_data() -> dict:
|
|
return pd.DataFrame(
|
|
{
|
|
"日期": ["2026-05-08", "2026-05-09"],
|
|
"开盘": [10.0, 10.2],
|
|
"收盘": [10.3, 10.4],
|
|
"最高": [10.5, 10.6],
|
|
"最低": [9.9, 10.0],
|
|
"成交量": [100, 120],
|
|
"成交额": [1030, 1248],
|
|
"振幅": [5.8, 5.9],
|
|
"涨跌幅": [3.0, 0.97],
|
|
"涨跌额": [0.3, 0.1],
|
|
"换手率": [1.0, 1.2],
|
|
}
|
|
)
|
|
|
|
|
|
def _make_bs_result(rows: list[list[str]]):
|
|
cursor = {"idx": -1}
|
|
|
|
def next_row():
|
|
cursor["idx"] += 1
|
|
return cursor["idx"] < len(rows)
|
|
|
|
def get_row_data():
|
|
return rows[cursor["idx"]]
|
|
|
|
return SimpleNamespace(error_code="0", error_msg="success", next=next_row, get_row_data=get_row_data)
|
|
|
|
|
|
@contextmanager
|
|
def _fake_bs_query(result):
|
|
yield result
|
|
|
|
|
|
def test_fetch_one_stock_converts_baostock_rows():
|
|
"""单股 BaoStock 读取应转成可入库行。"""
|
|
result = _make_bs_result([
|
|
["2026-05-08", "10.0", "10.5", "9.9", "10.3", "100", "1030", "10.0", "3.0", "1.0"],
|
|
["2026-05-09", "10.2", "10.6", "10.0", "10.4", "120", "1248", "10.3", "0.97", "1.2"],
|
|
])
|
|
|
|
with patch.object(daily, "bs_query", lambda *args, **kwargs: _fake_bs_query(result)):
|
|
rows = daily._fetch_one_stock("600000", "20260508", "20260509")
|
|
|
|
assert len(rows) == 2
|
|
assert rows[0]["code"] == "600000"
|
|
assert rows[0]["date"] == "2026-05-08"
|
|
assert rows[0]["open"] == 10.0
|
|
assert rows[0]["turnover"] == 1030
|
|
|
|
|
|
def test_fetch_one_stock_returns_empty_when_no_data():
|
|
"""BaoStock 返回空表时应返回空列表。"""
|
|
result = _make_bs_result([])
|
|
with patch.object(daily, "bs_query", lambda *args, **kwargs: _fake_bs_query(result)):
|
|
rows = daily._fetch_one_stock("600000", "20260508", "20260508")
|
|
|
|
assert rows == []
|
|
|
|
|
|
def test_fetch_one_stock_skips_rows_without_volume():
|
|
"""volume 为空的日线行不应进入 stock_daily。"""
|
|
result = _make_bs_result([
|
|
["2026-05-08", "10.0", "10.5", "9.9", "10.3", "", "1030", "10.0", "3.0", "1.0"],
|
|
["2026-05-09", "10.2", "10.6", "10.0", "10.4", "120", "1248", "10.3", "0.97", "1.2"],
|
|
])
|
|
|
|
with patch.object(daily, "bs_query", lambda *args, **kwargs: _fake_bs_query(result)):
|
|
rows = daily._fetch_one_stock("600000", "20260508", "20260509")
|
|
|
|
assert len(rows) == 1
|
|
assert rows[0]["date"] == "2026-05-09"
|
|
assert rows[0]["volume"] == 120
|
|
|
|
|
|
def test_fetch_daily_uses_baostock_fetcher():
|
|
"""fetch_daily 应走 BaoStock 抓取并落库。"""
|
|
fake_session = SimpleNamespace(execute=lambda *args, **kwargs: [], close=lambda: None)
|
|
hist_rows = [
|
|
["2026-05-08", "10.0", "10.5", "9.9", "10.3", "100", "1030", "10.0", "3.0", "1.0"],
|
|
]
|
|
result = _make_bs_result(hist_rows)
|
|
|
|
with patch.object(daily, "bs_query", lambda *args, **kwargs: _fake_bs_query(result)), \
|
|
patch.object(daily, "get_stock_codes", return_value=["600000"]), \
|
|
patch.object(daily, "get_trading_days", return_value=["2026-05-08"]), \
|
|
patch.object(daily, "get_session", return_value=fake_session), \
|
|
patch.object(daily, "_analyze_gaps", return_value={"600000": ["2026-05-08", "2026-05-08"]}), \
|
|
patch.object(daily, "batch_upsert") as mock_upsert:
|
|
daily.fetch_daily(start_date="20260508", end_date="20260508")
|
|
|
|
mock_upsert.assert_called_once()
|
|
|
|
|
|
def test_fetch_daily_does_not_mark_no_data_days_when_empty_rows():
|
|
"""fetch_daily 遇到空结果时不应写入停牌/无数据日期。"""
|
|
fake_session = SimpleNamespace(execute=lambda *args, **kwargs: [], close=lambda: None)
|
|
empty_result = _make_bs_result([])
|
|
|
|
with patch.object(daily, "bs_query", lambda *args, **kwargs: _fake_bs_query(empty_result)), \
|
|
patch.object(daily, "get_stock_codes", return_value=["600000"]), \
|
|
patch.object(daily, "get_trading_days", return_value=["2026-05-08", "2026-05-09"]), \
|
|
patch.object(daily, "get_session", return_value=fake_session), \
|
|
patch.object(daily, "_analyze_gaps", return_value={"600000": ["2026-05-08", "2026-05-09"]}), \
|
|
patch.object(daily, "batch_upsert") as mock_upsert:
|
|
daily.fetch_daily(start_date="20260508", end_date="20260509")
|
|
|
|
assert mock_upsert.call_count == 0
|
|
|
|
|
|
def test_fetch_daily_patches_missing_days_from_partial_range():
|
|
"""--daily 遇到范围结果缺日时，应按天补抓并写入完整数据。"""
|
|
fake_session = SimpleNamespace(execute=lambda *args, **kwargs: [], close=lambda: None)
|
|
range_rows = [
|
|
{"code": "600000", "date": "2026-05-08", "open": 10.0, "close": 10.3, "high": 10.5, "low": 9.9,
|
|
"volume": 100, "turnover": 1030, "amplitude": 5.8, "pct_change": 3.0, "change": 0.3, "turnover_rate": 1.0},
|
|
]
|
|
single_day_rows = [
|
|
{"code": "600000", "date": "2026-05-09", "open": 10.2, "close": 10.4, "high": 10.6, "low": 10.0,
|
|
"volume": 120, "turnover": 1248, "amplitude": 5.9, "pct_change": 0.97, "change": 0.1, "turnover_rate": 1.2},
|
|
]
|
|
|
|
def fake_fetch(code, gap_start, gap_end, *, retry=3, delay=0.1):
|
|
if gap_start == "20260508" and gap_end == "20260509":
|
|
return list(range_rows)
|
|
if gap_start == "20260509" and gap_end == "20260509":
|
|
return list(single_day_rows)
|
|
raise AssertionError((code, gap_start, gap_end))
|
|
|
|
with patch.object(daily, "get_stock_codes", return_value=["600000"]), \
|
|
patch.object(daily, "get_trading_days", return_value=["2026-05-08", "2026-05-09"]), \
|
|
patch.object(daily, "get_session", return_value=fake_session), \
|
|
patch.object(daily, "_analyze_gaps", return_value={"600000": ["2026-05-08", "2026-05-09"]}), \
|
|
patch.object(daily, "_fetch_one_stock_with_retry", side_effect=fake_fetch), \
|
|
patch.object(daily, "batch_upsert") as mock_upsert:
|
|
daily.fetch_daily(start_date="20260508", end_date="20260509")
|
|
|
|
mock_upsert.assert_called_once()
|
|
model_cls, rows, index_columns = mock_upsert.call_args.args
|
|
assert model_cls.__tablename__ == "stock_daily"
|
|
assert index_columns == ["code", "date"]
|
|
assert [row["date"] for row in rows] == ["2026-05-08", "2026-05-09"]
|
|
|
|
|
|
def test_fetch_daily_no_data_only_marks_partial_gap_days():
|
|
"""--daily-no-data-only 遇到部分行情时应补记缺失日期。"""
|
|
fake_session = SimpleNamespace(execute=lambda *args, **kwargs: [], close=lambda: None)
|
|
hist_rows = [
|
|
["2026-05-08", "10.0", "10.5", "9.9", "10.3", "100", "1030", "10.0", "3.0", "1.0"],
|
|
]
|
|
result = _make_bs_result(hist_rows)
|
|
|
|
with patch.object(daily, "bs_query", lambda *args, **kwargs: _fake_bs_query(result)), \
|
|
patch.object(daily, "get_stock_codes", return_value=["600000"]), \
|
|
patch.object(daily, "get_trading_days", return_value=["2026-05-08", "2026-05-09"]), \
|
|
patch.object(daily, "get_session", return_value=fake_session), \
|
|
patch.object(daily, "_analyze_gaps", return_value={"600000": ["2026-05-08", "2026-05-09"]}), \
|
|
patch.object(daily, "batch_upsert") as mock_upsert:
|
|
daily.fetch_daily(start_date="20260508", end_date="20260509", no_data_only=True)
|
|
|
|
assert mock_upsert.call_count == 1
|
|
model_cls, rows, index_columns = mock_upsert.call_args.args
|
|
assert model_cls.__tablename__ == "stock_no_data"
|
|
assert index_columns == ["code", "date"]
|
|
assert [row["date"] for row in rows] == ["2026-05-09"]
|
|
|
|
|
|
def test_fetch_daily_no_data_only_marks_empty_result_as_no_data():
|
|
"""--daily-no-data-only 遇到空结果时仍应写入停牌/无数据日期。"""
|
|
fake_session = SimpleNamespace(execute=lambda *args, **kwargs: [], close=lambda: None)
|
|
empty_result = _make_bs_result([])
|
|
|
|
with patch.object(daily, "bs_query", lambda *args, **kwargs: _fake_bs_query(empty_result)), \
|
|
patch.object(daily, "get_stock_codes", return_value=["600000"]), \
|
|
patch.object(daily, "get_trading_days", return_value=["2026-05-08", "2026-05-09"]), \
|
|
patch.object(daily, "get_session", return_value=fake_session), \
|
|
patch.object(daily, "_analyze_gaps", return_value={"600000": ["2026-05-08", "2026-05-09"]}), \
|
|
patch.object(daily, "batch_upsert") as mock_upsert:
|
|
daily.fetch_daily(start_date="20260508", end_date="20260509", no_data_only=True)
|
|
|
|
assert mock_upsert.call_count == 1
|
|
model_cls, rows, index_columns = mock_upsert.call_args.args
|
|
assert model_cls.__tablename__ == "stock_no_data"
|
|
assert index_columns == ["code", "date"]
|
|
assert [row["date"] for row in rows] == ["2026-05-08", "2026-05-09"]
|
|
|
|
|
|
def test_fetch_daily_skips_unsupported_non_a_share_codes():
|
|
"""fetch_daily 应在请求前过滤非沪深 A 股代码。"""
|
|
fake_session = SimpleNamespace(execute=lambda *args, **kwargs: [], close=lambda: None)
|
|
|
|
with patch.object(daily, "get_stock_codes", return_value=["600000", "920200"]), \
|
|
patch.object(daily, "get_trading_days", return_value=["2026-05-08"]), \
|
|
patch.object(daily, "get_session", return_value=fake_session), \
|
|
patch.object(daily, "_analyze_gaps", return_value={"600000": ["2026-05-08", "2026-05-08"]}) as mock_analyze, \
|
|
patch.object(daily, "_fetch_one_stock_with_retry", return_value=[]) as mock_fetch, \
|
|
patch.object(daily, "batch_upsert"):
|
|
daily.fetch_daily(start_date="20260508", end_date="20260508")
|
|
|
|
mock_analyze.assert_called_once()
|
|
assert mock_analyze.call_args.args[0] == ["600000"]
|
|
mock_fetch.assert_called_once_with("600000", "20260508", "20260508", retry=3, delay=0.2)
|
|
|
|
|
|
def test_fetch_daily_retries_transient_baostock_failure():
|
|
"""fetch_daily 遇到临时断连时应重试。"""
|
|
fake_session = SimpleNamespace(execute=lambda *args, **kwargs: [], close=lambda: None)
|
|
hist_rows = [
|
|
["2026-05-08", "10.0", "10.5", "9.9", "10.3", "100", "1030", "10.0", "3.0", "1.0"],
|
|
]
|
|
result = _make_bs_result(hist_rows)
|
|
calls = {"count": 0}
|
|
|
|
def flaky_query(*args, **kwargs):
|
|
calls["count"] += 1
|
|
if calls["count"] == 1:
|
|
raise ConnectionError("boom")
|
|
return _fake_bs_query(result)
|
|
|
|
with patch.object(daily, "bs_query", flaky_query), \
|
|
patch.object(daily, "get_fetch_config", return_value={"delay": 0, "retry": 2}), \
|
|
patch.object(daily, "get_stock_codes", return_value=["600000"]), \
|
|
patch.object(daily, "get_trading_days", return_value=["2026-05-08"]), \
|
|
patch.object(daily, "get_session", return_value=fake_session), \
|
|
patch.object(daily, "_analyze_gaps", return_value={"600000": ["2026-05-08", "2026-05-08"]}), \
|
|
patch.object(daily, "batch_upsert") as mock_upsert, \
|
|
patch.object(daily.time, "sleep", return_value=None):
|
|
daily.fetch_daily(start_date="20260508", end_date="20260508")
|
|
|
|
assert calls["count"] == 2
|
|
mock_upsert.assert_called_once()
|
|
|
|
|
|
def test_analyze_gaps_marks_incomplete_volume_as_gap():
|
|
"""已存在但关键字段为空时，应视为缺口并重抓。"""
|
|
class FakeSession:
|
|
def __init__(self):
|
|
self.calls = 0
|
|
|
|
def execute(self, *args, **kwargs):
|
|
self.calls += 1
|
|
return [("600000", "2026-05", 2, 1)]
|
|
|
|
def close(self):
|
|
return None
|
|
|
|
fake_session = FakeSession()
|
|
|
|
with patch.object(daily, "get_ipo_dates", return_value={"600000": "2026-01-01"}), \
|
|
patch.object(daily, "_get_no_data_dates", return_value={}), \
|
|
patch.object(daily, "get_session", return_value=fake_session):
|
|
gaps = daily._analyze_gaps(
|
|
["600000"],
|
|
"20260508",
|
|
"20260509",
|
|
["2026-05-08", "2026-05-09"],
|
|
)
|
|
|
|
assert gaps["600000"] == ["2026-05-08", "2026-05-09"]
|
|
assert fake_session.calls == 1
|
|
|
|
|
|
def test_analyze_gaps_detects_missing_middle_trading_day():
|
|
"""已有头尾数据但中间缺交易日时，也应纳入缺口。"""
|
|
class FakeSession:
|
|
def __init__(self):
|
|
self.calls = 0
|
|
|
|
def execute(self, *args, **kwargs):
|
|
self.calls += 1
|
|
if self.calls == 1:
|
|
return [("600000", "2026-05", 2, 0)]
|
|
return [("2026-05-08",), ("2026-05-11",)]
|
|
|
|
def close(self):
|
|
return None
|
|
|
|
fake_session = FakeSession()
|
|
|
|
with patch.object(daily, "get_ipo_dates", return_value={"600000": "2026-01-01"}), \
|
|
patch.object(daily, "_get_no_data_dates", return_value={}), \
|
|
patch.object(daily, "get_session", return_value=fake_session):
|
|
gaps = daily._analyze_gaps(
|
|
["600000"],
|
|
"20260508",
|
|
"20260511",
|
|
["2026-05-08", "2026-05-09", "2026-05-11"],
|
|
)
|
|
|
|
assert gaps["600000"] == ["2026-05-09", "2026-05-09"]
|
|
assert fake_session.calls == 2
|