50 lines
2.5 KiB
Python
50 lines
2.5 KiB
Python
# tests/data/test_datafeed.py
|
|
import time
|
|
import pandas as pd
|
|
import pytest
|
|
from unittest.mock import patch
|
|
from sanguo_data.datafeed import fetch_with_fallback, _fetch_baostock_with_timeout
|
|
|
|
def test_fetch_with_fallback_uses_second_when_first_fails():
|
|
df_good = pd.DataFrame({"date": ["2026-01-01"], "open": [10.0], "high": [11.0], "low": [9.0], "close": [10.5], "volume": [1000], "amount": [10000]})
|
|
with patch("sanguo_data.datafeed._fetch_eastmoney", side_effect=Exception("limit")), \
|
|
patch("sanguo_data.datafeed._fetch_baostock_with_timeout", return_value=df_good):
|
|
out = fetch_with_fallback("600000", "2026-01-01", "2026-01-02", ["eastmoney", "baostock"])
|
|
assert len(out) == 1
|
|
assert out["date"].iloc[0] == "2026-01-01"
|
|
|
|
def test_baostock_timeout_does_not_hang():
|
|
"""v1 卡死坑修复验证：超时必须返回，不能无限挂起"""
|
|
import sanguo_data.datafeed as df_module
|
|
|
|
# Save original worker
|
|
original_worker = df_module._baostock_worker
|
|
|
|
try:
|
|
# Replace with hanging worker (module-level function can be pickled)
|
|
df_module._baostock_worker = df_module._hanging_worker_for_test
|
|
|
|
start = time.time()
|
|
with pytest.raises(TimeoutError):
|
|
df_module._fetch_baostock_with_timeout("600000", "2026-01-01", "2026-01-02", timeout=2)
|
|
elapsed = time.time() - start
|
|
assert elapsed < 5, f"Timeout test took {elapsed:.2f}s, expected <5s"
|
|
finally:
|
|
# Restore original worker
|
|
df_module._baostock_worker = original_worker
|
|
|
|
def test_fetch_with_fallback_all_sources_fail():
|
|
"""所有源都失败时应该抛出异常"""
|
|
with patch("sanguo_data.datafeed._fetch_eastmoney", side_effect=Exception("em failed")), \
|
|
patch("sanguo_data.datafeed._fetch_baostock_with_timeout", side_effect=Exception("bs failed")):
|
|
with pytest.raises(RuntimeError, match="all sources failed"):
|
|
fetch_with_fallback("600000", "2026-01-01", "2026-01-02", ["eastmoney", "baostock"])
|
|
|
|
def test_fetch_with_fallback_first_succeeds():
|
|
"""第一个源成功时直接返回"""
|
|
df_good = pd.DataFrame({"date": ["2026-01-01"], "open": [10.0], "high": [11.0], "low": [9.0], "close": [10.5], "volume": [1000], "amount": [10000]})
|
|
with patch("sanguo_data.datafeed._fetch_eastmoney", return_value=df_good):
|
|
out = fetch_with_fallback("600000", "2026-01-01", "2026-01-02", ["eastmoney", "baostock"])
|
|
assert len(out) == 1
|
|
assert out["date"].iloc[0] == "2026-01-01"
|