diff --git a/config/data_platform.yaml b/config/data_platform.yaml new file mode 100644 index 0000000..3269559 --- /dev/null +++ b/config/data_platform.yaml @@ -0,0 +1,34 @@ +# config/data_platform.yaml +data_paths: + daily_dir: /volume1/stock/A股数据/日线数据/daily + minute_15_dir: /volume1/stock/minute_kline/15min + vnpy_db: /volume1/stock/sanguo_vnpy/data/quant_trading.db + stock_list: /volume1/stock/A股数据/stock_info/stock_basic_info_raw_20260326_113530.csv + +data_sources: + daily: + - name: eastmoney + enabled: true + interval: 4.0 + - name: baostock + enabled: true + interval: 0.0 + timeout: 30 + - name: tencent + enabled: true + interval: 0.0 + minute_15: + - name: eastmoney + enabled: true + interval: 4.0 + +validation: + price_positive: true + ohlc_consistency: true + no_future_dates: true + +performance: + request_interval: 0.3 + max_retries: 3 + fail_window: 100 + fail_threshold: 0.8 diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..c7b23ec --- /dev/null +++ b/pytest.ini @@ -0,0 +1,3 @@ +[pytest] +pythonpath = . +testpaths = tests diff --git a/sanguo_data/__init__.py b/sanguo_data/__init__.py index e69de29..cc749f9 100644 --- a/sanguo_data/__init__.py +++ b/sanguo_data/__init__.py @@ -0,0 +1,3 @@ +# sanguo_data/__init__.py +from .config import DataConfig, load_config +__all__ = ["DataConfig", "load_config"] diff --git a/sanguo_data/config.py b/sanguo_data/config.py new file mode 100644 index 0000000..5259b7d --- /dev/null +++ b/sanguo_data/config.py @@ -0,0 +1,29 @@ +# sanguo_data/config.py +from dataclasses import dataclass +import yaml + +@dataclass(frozen=True) +class DataConfig: + data_paths: dict + data_sources: dict + validation: dict + performance: dict + +def load_config(path: str) -> DataConfig: + try: + with open(path, "r", encoding="utf-8") as f: + raw = yaml.safe_load(f) + except FileNotFoundError: + raise FileNotFoundError(f"配置文件不存在: {path}") + except yaml.YAMLError as e: + raise ValueError(f"YAML解析失败: {e}") + + if not raw: + raise ValueError(f"配置文件为空: {path}") + + return DataConfig( + data_paths=raw.get("data_paths", {}), + data_sources=raw.get("data_sources", {}), + validation=raw.get("validation", {}), + performance=raw.get("performance", {}), + ) diff --git a/sanguo_data/datafeed.py b/sanguo_data/datafeed.py new file mode 100644 index 0000000..a3f833b --- /dev/null +++ b/sanguo_data/datafeed.py @@ -0,0 +1,228 @@ +# sanguo_data/datafeed.py +import pandas as pd +import urllib.request +import json +import time +import logging +from datetime import datetime, timedelta +from multiprocessing import Process, Queue +from typing import Optional +from sanguo_data.config import DataConfig + +logger = logging.getLogger(__name__) + +def fetch_with_fallback(symbol, start, end, sources: list[str]) -> pd.DataFrame: + fetchers = { + "eastmoney": _fetch_eastmoney, + "baostock": lambda s, st, e: _fetch_baostock_with_timeout(s, st, e, timeout=30), + "tencent": _fetch_tencent, + } + last_err = None + for name in sources: + try: + df = fetchers[name](symbol, start, end) + if df is not None and len(df) > 0: + return df + except Exception as e: + last_err = e + continue + raise RuntimeError(f"all sources failed: {last_err}") + +def fetch_daily(symbol, start, end, cfg: DataConfig) -> pd.DataFrame: + sources = [s["name"] for s in cfg.data_sources.get("daily", []) if s.get("enabled", True)] + return fetch_with_fallback(symbol, start, end, sources) + +# Worker function for multiprocessing (must be at module level to be picklable) +def _baostock_worker(symbol, start, end, result_queue): + try: + result = _fetch_baostock_raw(symbol, start, end) + result_queue.put(result) + except Exception as e: + result_queue.put(e) + +# Test helper for timeout testing (simulates hanging BaoStock call) +def _hanging_worker_for_test(symbol, start, end, result_queue): + """Test helper: simulates a hanging BaoStock call (60s sleep)""" + import time as _time + _time.sleep(60) + result_queue.put(pd.DataFrame({"date": ["2026-01-01"], "open": [10.0]})) + +def _fetch_baostock_with_timeout(symbol, start, end, timeout=30): + """子进程隔离 BaoStock(修复 v1 无超时卡死坑)""" + result_queue = Queue() + p = Process(target=_baostock_worker, args=(symbol, start, end, result_queue)) + p.start() + p.join(timeout) + if p.is_alive(): + p.terminate() + p.join() + raise TimeoutError(f"baostock timeout after {timeout}s") + + res = result_queue.get() + if isinstance(res, Exception): + raise res + return res + +def _get_em_secid(code: str) -> str: + if code.startswith(("60", "68", "51")): + return f"1.{code}" + return f"0.{code}" + +def _parse_em_klines(klines: list) -> Optional[pd.DataFrame]: + """解析东方财富K线数据(日线和15min通用)""" + if not klines: + return None + rows = [] + for line in klines: + parts = line.split(",") + if len(parts) < 7: + continue + rows.append({ + "date": parts[0], + "open": float(parts[1]), + "close": float(parts[2]), + "high": float(parts[3]), + "low": float(parts[4]), + "volume": float(parts[5]), + "amount": float(parts[6]), + }) + if not rows: + return None + return pd.DataFrame(rows) + +def _fetch_baostock_raw(symbol: str, start_date: str, end_date: str) -> Optional[pd.DataFrame]: + """BaoStock日线:全量历史,无反爬,amount真实,T+1延迟 + + Copied from v1 data_platform/daily_all_update.py:fetch_baostock_daily (lines 242-270) + """ + try: + import baostock as bs + except ImportError: + return None + + # 转换代码格式:600000 -> sh.600000 + code = symbol.replace("SH", "").replace("SZ", "").replace("sh", "").replace("sz", "") + if code.startswith(("60", "68", "51")): + bs_code = f"sh.{code}" + else: + bs_code = f"sz.{code}" + + try: + rs = bs.query_history_k_data_plus( + bs_code, + "date,open,high,low,close,volume,amount", + start_date=start_date.replace("-", ""), + end_date=end_date.replace("-", ""), + frequency="d", + adjustflag="2", + ) + rows = [] + while (rs.error_code == "0") and rs.next(): + rows.append(rs.get_row_data()) + if not rows: + return None + df = pd.DataFrame(rows, columns=["date", "open", "high", "low", "close", "volume", "amount"]) + for c in ["open", "high", "low", "close", "volume", "amount"]: + df[c] = pd.to_numeric(df[c], errors="coerce") + df = df.dropna(subset=["close"]) + if df.empty: + return None + return df + except Exception as e: + logger.debug("BaoStock日线失败 %s: %s", symbol, e) + return None + +def _fetch_eastmoney(symbol: str, start_date: str, end_date: str) -> Optional[pd.DataFrame]: + """东方财富日线:当天实时,amount真实,4s限频 + + Copied from v1 data_platform/daily_all_update.py:fetch_eastmoney_daily (lines 339-374) + """ + try: + import requests as _requests + except ImportError: + return None + + code = symbol.replace("SH", "").replace("SZ", "").replace("sh", "").replace("sz", "") + secid = _get_em_secid(code) + ts = str(int(time.time() * 1000)) + url = ( + f"https://push2his.eastmoney.com/api/qt/stock/kline/get?" + f"secid={secid}&klt=101&fqt=1&" + f"beg={start_date.replace('-', '')}&end={end_date.replace('-', '')}&" + f"fields1=f1,f2,f3,f4,f5,f6,f7,f8&" + f"fields2=f51,f52,f53,f54,f55,f56,f57,f58,f59,f60,f61&" + f"ut=b2884a393a59ad64002292a3e90d46a5&lmt=10000&" + f"cb=jQuery_em_{ts}&_={ts}" + ) + + headers_em = { + "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36", + "Referer": "https://quote.eastmoney.com/", + "Accept": "*/*", + "Accept-Language": "zh-CN,zh;q=0.9", + } + + session = _requests.Session() + session.trust_env = False + try: + r = session.get(url, headers=headers_em, timeout=15, verify=False) + if r.status_code != 200: + return None + text = r.text + data = json.loads(text[text.index("(") + 1:text.rindex(")")]) + if data.get("rc") != 0: + return None + klines = data.get("data", {}).get("klines", []) + df = _parse_em_klines(klines) + if df is None: + return None + df["date"] = pd.to_datetime(df["date"]).dt.strftime("%Y-%m-%d") + mask = (df["date"] >= start_date) & (df["date"] <= end_date) + result = df.loc[mask, ["date", "open", "high", "low", "close", "volume", "amount"]] + return result if not result.empty else None + except Exception as e: + logger.debug("东方财富日线失败 %s: %s", symbol, e) + return None + +def _fetch_tencent(symbol: str, start_date: str, end_date: str) -> Optional[pd.DataFrame]: + """腾讯日线:amount有时为0 + + Copied from v1 data_platform/fallback.py:_fetch_tencent_daily (lines 66-104) + """ + code = symbol.replace("SH", "").replace("SZ", "").replace("sh", "").replace("sz", "") + if code.startswith(("6", "5", "1")): + prefix = "sh" + else: + prefix = "sz" + tq_symbol = f"{prefix}{code}" + + days = (datetime.strptime(end_date, "%Y-%m-%d") - datetime.strptime(start_date, "%Y-%m-%d")).days + 10 + url = f"https://web.ifzq.gtimg.cn/appstock/app/fqkline/get?param={tq_symbol},day,{start_date},,{days}," + + opener = urllib.request.build_opener(urllib.request.ProxyHandler({})) + try: + req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"}) + with opener.open(req, timeout=10) as r: + resp = json.loads(r.read()) + d = resp.get("data") + if not isinstance(d, dict): + return None + klines = d.get(tq_symbol, {}).get("day", []) + if not klines: + return None + df = pd.DataFrame(klines) + ncols = len(df.columns) + if ncols >= 7: + df.columns = ["date", "open", "close", "high", "low", "volume", "amount"][:ncols] + else: + df.columns = ["date", "open", "close", "high", "low", "volume"][:ncols] + if "amount" not in df.columns: + df["amount"] = 0.0 + for c in ["open", "close", "high", "low", "volume", "amount"]: + df[c] = pd.to_numeric(df[c], errors="coerce").fillna(0) + df["date"] = pd.to_datetime(df["date"]).dt.strftime("%Y-%m-%d") + mask = (df["date"] >= start_date) & (df["date"] <= end_date) + return df.loc[mask, ["date", "open", "high", "low", "close", "volume", "amount"]] + except Exception as e: + logger.debug("腾讯日线失败 %s: %s", symbol, e) + return None \ No newline at end of file diff --git a/sanguo_data/datareader.py b/sanguo_data/datareader.py new file mode 100644 index 0000000..bdb4718 --- /dev/null +++ b/sanguo_data/datareader.py @@ -0,0 +1,67 @@ +import sys +import os +from pathlib import Path + +# Add real vnpy source code to sys.path +_VNPY_SRC = os.path.join(os.path.dirname(__file__), "..", "..", "vnpy_v4.4.0") +_VNPY_SRC = os.path.abspath(_VNPY_SRC) +if _VNPY_SRC not in sys.path: + sys.path.insert(0, _VNPY_SRC) + +import pandas as pd +from datetime import datetime +from vnpy.trader.object import BarData +from vnpy.trader.constant import Exchange, Interval +from vnpy.trader.database import get_database + +def read_parquet_daily(symbol: str, start: str, end: str, cfg) -> list[BarData]: + daily_dir = Path(cfg.data_paths["daily_dir"]) + start_dt = datetime.strptime(start, "%Y-%m-%d") + end_dt = datetime.strptime(end, "%Y-%m-%d") + bars: list[BarData] = [] + for year in range(start_dt.year, end_dt.year + 1): + f = daily_dir / str(year) / f"{symbol}.parquet" + if not f.exists(): + continue + df = pd.read_parquet(f) + for _, row in df.iterrows(): + d = pd.to_datetime(row["date"]) + if start_dt <= d <= end_dt: + bars.append(_row_to_bar(symbol, row, Interval.DAILY)) + return bars + +def _row_to_bar(symbol: str, row, interval: Interval) -> BarData: + return BarData( + symbol=symbol, + exchange=guess_exchange(symbol), # Task 4 已改为 guess_exchange + datetime=pd.to_datetime(row["date"]).to_pydatetime(), + interval=interval, + open_price=float(row["open"]), + high_price=float(row["high"]), + low_price=float(row["low"]), + close_price=float(row["close"]), + volume=float(row["volume"]), + gateway_name="DATA", + ) + + +def guess_exchange(symbol: str) -> Exchange: + """按代码前缀判断交易所:6/68/5x→SSE,0/3/15x→SZSE""" + if symbol.startswith(("60", "68", "51", "56", "58")): + return Exchange.SSE + if symbol.startswith(("00", "30", "15")): + return Exchange.SZSE + return Exchange.SSE + + +def read_db_daily(symbol: str, start: str, end: str, cfg) -> list[BarData]: + db = get_database() + start_dt = datetime.strptime(start, "%Y-%m-%d") + end_dt = datetime.strptime(end, "%Y-%m-%d") + return db.load_bar_data( + symbol=symbol, + exchange=guess_exchange(symbol), + interval=Interval.DAILY, + start=start_dt, + end=end_dt, + ) diff --git a/sanguo_data/datawriter.py b/sanguo_data/datawriter.py new file mode 100644 index 0000000..237abb0 --- /dev/null +++ b/sanguo_data/datawriter.py @@ -0,0 +1,34 @@ +# sanguo_data/datawriter.py +import os +import pandas as pd +from pathlib import Path +from vnpy.trader.object import BarData +from vnpy.trader.constant import Interval +from sanguo_data.datareader import _row_to_bar +from sanguo_data.config import DataConfig + +def atomic_write_parquet(path: str, df: pd.DataFrame) -> None: + p = Path(path) + p.parent.mkdir(parents=True, exist_ok=True) + tmp = str(p) + ".tmp" + df.to_parquet(tmp) + os.replace(tmp, str(p)) # 原子替换 + +def write_daily(symbol: str, df: pd.DataFrame, cfg: DataConfig) -> None: + # 1) parquet 增量合并(按年分区,去重保留最新) + for year, group in df.groupby(df["date"].str[:4]): + f = Path(cfg.data_paths["daily_dir"]) / year / f"{symbol}.parquet" + if f.exists(): + old = pd.read_parquet(f) + combined = pd.concat([old, group]).drop_duplicates("date", keep="last") + else: + combined = group + atomic_write_parquet(str(f), combined) + # 2) vnpy SQLite + bars = [_row_to_bar(symbol, row, Interval.DAILY) for _, row in df.iterrows()] + _save_to_vnpy_db(bars, cfg) + +def _save_to_vnpy_db(bars: list[BarData], cfg: DataConfig) -> None: + from vnpy.trader.database import get_database + db = get_database() + db.save_bar_data(bars) # spike 验证签名 diff --git a/sanguo_data/scheduler.py b/sanguo_data/scheduler.py new file mode 100644 index 0000000..73cabc2 --- /dev/null +++ b/sanguo_data/scheduler.py @@ -0,0 +1,57 @@ +# sanguo_data/scheduler.py +import json +import time +from dataclasses import dataclass, field +from pathlib import Path +from sanguo_data.config import DataConfig +from sanguo_data.datafeed import fetch_daily +from sanguo_data.validator import validate_daily +from sanguo_data.datawriter import write_daily + +@dataclass +class UpdateReport: + total: int = 0 + success: int = 0 + failed: int = 0 + skipped: int = 0 + failures: list = field(default_factory=list) + +def run_daily_update(cfg: DataConfig, symbols: list[str] | None = None) -> UpdateReport: + progress_path = Path(cfg.data_paths.get("progress_file", "progress.json")) + progress = json.loads(progress_path.read_text()) if progress_path.exists() else {} + symbols = symbols or _load_stock_list(cfg) + report = UpdateReport(total=len(symbols)) + fail_window = cfg.performance.get("fail_window", 100) + fail_threshold = cfg.performance.get("fail_threshold", 0.8) + + for sym in symbols: + if progress.get(sym) == "done": + report.skipped += 1 + continue + try: + df = fetch_daily(sym, _last_date(sym, cfg), _today(), cfg) + df = validate_daily(df) + if len(df) > 0: + write_daily(sym, df, cfg) + progress[sym] = "done" + progress_path.write_text(json.dumps(progress, ensure_ascii=False)) + report.success += 1 + except Exception as e: + report.failed += 1 + report.failures.append({"symbol": sym, "error": str(e)}) + checked = report.success + report.failed + if checked >= fail_window and report.failed / max(checked, 1) > fail_threshold: + report.failures.append({"error": "FAIL_THRESHOLD_REACHED, abort"}) + break + time.sleep(cfg.performance.get("request_interval", 0.3)) + return report + +def _load_stock_list(cfg: DataConfig) -> list[str]: + """从 v1 data_platform/daily_all_update.py copy 全市场股票列表读取""" + raise NotImplementedError("copy from v1") + +def _last_date(symbol: str, cfg: DataConfig) -> str: + return "2020-01-01" # 简化,实际读 parquet 最后日期 + +def _today() -> str: + return "2026-07-05" diff --git a/sanguo_data/validator.py b/sanguo_data/validator.py new file mode 100644 index 0000000..6f6ffc5 --- /dev/null +++ b/sanguo_data/validator.py @@ -0,0 +1,132 @@ +#!/usr/bin/env python3 +"""数据校验层 - V1 7条fatal规则""" +import pandas as pd +from datetime import datetime +from typing import List, Tuple + + +class ValidationResult: + def __init__(self): + self.passed = True + self.fatal_errors: List[str] = [] + self.warnings: List[str] = [] + self.checked_rows = 0 + self.failed_rows = 0 + + def __repr__(self): + return (f"ValidationResult(passed={self.passed}, " + f"fatal={len(self.fatal_errors)}, warnings={len(self.warnings)}, " + f"rows={self.checked_rows}, failed={self.failed_rows})") + + def to_dict(self): + return { + "passed": self.passed, + "fatal_errors": self.fatal_errors, + "warnings": self.warnings, + "checked_rows": self.checked_rows, + "failed_rows": self.failed_rows, + } + + +class DataValidator: + """数据校验器 - V1 7条fatal规则""" + + def validate(self, df: pd.DataFrame, data_type: str = "daily") -> ValidationResult: + result = ValidationResult() + if df is None or df.empty: + result.fatal_errors.append("数据为空") + result.passed = False + return result + result.checked_rows = len(df) + + if data_type == "daily": + self._validate_daily(df, result) + elif data_type == "realtime": + self._validate_realtime(df, result) + return result + + def validate_realtime_dict(self, data: dict) -> ValidationResult: + """校验单条实时行情""" + result = ValidationResult() + result.checked_rows = 1 + errors = [] + # R1: 价格>0 + if not data or data.get("current", 0) <= 0: + errors.append("R1: current价格<=0") + if data.get("prev_close", 0) <= 0: + errors.append("R1: prev_close<=0") + # R7: 必须携带source和fetched_at + if not data.get("source"): + errors.append("R7: 缺少source字段") + if not data.get("fetched_at"): + errors.append("R7: 缺少fetched_at字段") + if errors: + result.fatal_errors = errors + result.passed = False + result.failed_rows = 1 + return result + + def _validate_daily(self, df: pd.DataFrame, result: ValidationResult): + today = datetime.now().strftime("%Y-%m-%d") + + for idx, row in df.iterrows(): + row_errors = [] + # D1: 价格>0 + for col in ["close", "open", "high", "low"]: + val = row.get(col, 0) + if pd.isna(val) or float(val) <= 0: + row_errors.append(f"D1: {col}<=0 (row {idx})") + break + + # D2: OHLC一致性 + o, h, l, c = float(row.get("open", 0)), float(row.get("high", 0)), \ + float(row.get("low", 0)), float(row.get("close", 0)) + if o > 0 and c > 0: + if h < max(o, c) or l > min(o, c): + row_errors.append(f"D2: OHLC不一致 (row {idx}, o={o} h={h} l={l} c={c})") + + # D3: volume >= 0 + vol = row.get("volume", 0) + if pd.notna(vol) and float(vol) < 0: + row_errors.append(f"D3: volume<0 (row {idx})") + + # D7: 非未来日期 + dt = str(row.get("date", row.get("datetime", "")))[:10] + if dt > today: + row_errors.append(f"D7: 未来日期 {dt} (row {idx})") + + if row_errors: + result.fatal_errors.extend(row_errors) + result.failed_rows += 1 + + # D6: 日期不重复 (check after all rows) + date_col = "date" if "date" in df.columns else "datetime" + if date_col in df.columns: + dupes = df[df.duplicated(subset=[date_col], keep=False)] + if not dupes.empty and len(df) > 1: + result.fatal_errors.append(f"D6: {len(dupes)}条重复日期") + + if result.fatal_errors: + result.passed = False + + +# 适配层,不改 v1 校验逻辑 +def validate_daily(df): + """对外统一接口,委托 v1 校验规则 + + Args: + df: 输入日线DataFrame + + Returns: + 过滤后的DataFrame(仅包含通过校验的行) + """ + validator = DataValidator() + valid_rows = [] + + for idx in range(len(df)): + row_df = df.iloc[idx:idx+1] + result = validator.validate(row_df, data_type="daily") + if result.passed: + valid_rows.append(idx) + + return df.iloc[valid_rows] diff --git a/tests/data/__init__.py b/tests/data/__init__.py new file mode 100644 index 0000000..29e9380 --- /dev/null +++ b/tests/data/__init__.py @@ -0,0 +1 @@ +# tests/data/__init__.py diff --git a/tests/data/conftest.py b/tests/data/conftest.py new file mode 100644 index 0000000..bafac82 --- /dev/null +++ b/tests/data/conftest.py @@ -0,0 +1,30 @@ +import sys +import os + +# Add real vnpy source code to sys.path +_VNPY_SRC = os.path.join(os.path.dirname(__file__), "..", "..", "vnpy_v4.4.0") +_VNPY_SRC = os.path.abspath(_VNPY_SRC) +if _VNPY_SRC not in sys.path: + sys.path.insert(0, _VNPY_SRC) + +from pathlib import Path + +import pandas as pd +import pytest + +@pytest.fixture +def good_daily_df(): + return pd.DataFrame({ + "date": ["2026-01-01", "2026-01-02"], + "open": [10.0, 11.0], "high": [10.5, 11.5], + "low": [9.8, 10.8], "close": [10.2, 11.2], + "volume": [10000, 12000], + }) + +@pytest.fixture +def bad_daily_df(): + return pd.DataFrame({ + "date": ["2026-01-01"], + "open": [0.0], "high": [0.0], "low": [0.0], "close": [0.0], + "volume": [100], + }) \ No newline at end of file diff --git a/tests/data/test_config.py b/tests/data/test_config.py new file mode 100644 index 0000000..d66f1e2 --- /dev/null +++ b/tests/data/test_config.py @@ -0,0 +1,32 @@ +# tests/data/test_config.py +import pytest +from sanguo_data.config import load_config, DataConfig + +def test_load_config_returns_dataconfig(tmp_path): + yaml_content = """ +data_paths: + daily_dir: /tmp/daily + minute_15_dir: /tmp/15min + vnpy_db: /tmp/quant.db + stock_list: /tmp/stock.csv +data_sources: + daily: + - name: eastmoney + enabled: true + interval: 4.0 +validation: + price_positive: true +performance: + max_retries: 3 +""" + p = tmp_path / "config.yaml" + p.write_text(yaml_content) + cfg = load_config(str(p)) + assert isinstance(cfg, DataConfig) + assert cfg.data_paths["daily_dir"] == "/tmp/daily" + assert cfg.data_sources["daily"][0]["name"] == "eastmoney" + assert cfg.performance["max_retries"] == 3 + +def test_load_config_raises_on_missing_file(): + with pytest.raises(FileNotFoundError, match="配置文件不存在"): + load_config("/nonexistent/path/config.yaml") diff --git a/tests/data/test_datafeed.py b/tests/data/test_datafeed.py new file mode 100644 index 0000000..fc69d73 --- /dev/null +++ b/tests/data/test_datafeed.py @@ -0,0 +1,49 @@ +# 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" diff --git a/tests/data/test_datareader.py b/tests/data/test_datareader.py new file mode 100644 index 0000000..186baa2 --- /dev/null +++ b/tests/data/test_datareader.py @@ -0,0 +1,32 @@ +import pandas as pd +from vnpy.trader.constant import Exchange, Interval +from sanguo_data.config import DataConfig +from sanguo_data.datareader import read_parquet_daily, guess_exchange + +def test_read_parquet_daily_returns_bardata(tmp_path): + year_dir = tmp_path / "2026" + year_dir.mkdir() + df = pd.DataFrame({ + "date": ["2026-01-05", "2026-01-06"], + "open": [10.0, 11.0], "high": [10.5, 11.5], + "low": [9.8, 10.8], "close": [10.2, 11.2], + "volume": [10000, 12000], + }) + df.to_parquet(year_dir / "600000.parquet") + + cfg = DataConfig( + data_paths={"daily_dir": str(tmp_path)}, + data_sources={}, validation={}, performance={}, + ) + bars = read_parquet_daily("600000", "2026-01-01", "2026-12-31", cfg) + assert len(bars) == 2 + assert bars[0].symbol == "600000" + assert bars[0].open_price == 10.0 + + +def test_guess_exchange_sh(): + assert guess_exchange("600000").value == "SSE" + + +def test_guess_exchange_sz(): + assert guess_exchange("000001").value == "SZSE" diff --git a/tests/data/test_datawriter.py b/tests/data/test_datawriter.py new file mode 100644 index 0000000..58399e2 --- /dev/null +++ b/tests/data/test_datawriter.py @@ -0,0 +1,24 @@ +# tests/data/test_datawriter.py +import pandas as pd +from sanguo_data.config import DataConfig +from sanguo_data.datawriter import write_daily, atomic_write_parquet + +def test_atomic_write_parquet(tmp_path): + f = tmp_path / "2026" / "600000.parquet" + df = pd.DataFrame({"date": ["2026-01-01"], "open": [10.0]}) + atomic_write_parquet(str(f), df) + assert f.exists() + assert not list(tmp_path.glob("*.tmp")) + +def test_write_daily_writes_parquet_and_db(tmp_path, monkeypatch): + cfg = DataConfig( + data_paths={"daily_dir": str(tmp_path / "daily"), "vnpy_db": str(tmp_path / "q.db")}, + data_sources={}, validation={}, performance={}, + ) + df = pd.DataFrame({"date": ["2026-01-01"], "open": [10.0], "high": [10.0], + "low": [10.0], "close": [10.0], "volume": [100]}) + called = {} + monkeypatch.setattr("sanguo_data.datawriter._save_to_vnpy_db", lambda bars, cfg: called.setdefault("bars", bars)) + write_daily("600000", df, cfg) + assert (tmp_path / "daily" / "2026" / "600000.parquet").exists() + assert len(called["bars"]) == 1 diff --git a/tests/data/test_scheduler.py b/tests/data/test_scheduler.py new file mode 100644 index 0000000..69dd367 --- /dev/null +++ b/tests/data/test_scheduler.py @@ -0,0 +1,24 @@ +# tests/data/test_scheduler.py +import json +import pandas as pd +from unittest.mock import patch +from sanguo_data.config import DataConfig +from sanguo_data.scheduler import run_daily_update, UpdateReport + +def test_run_daily_update_skips_completed_on_resume(tmp_path): + progress_file = tmp_path / "progress.json" + progress_file.write_text('{"600000": "done"}') + cfg = DataConfig( + data_paths={"daily_dir": str(tmp_path), "vnpy_db": str(tmp_path / "q.db"), + "progress_file": str(progress_file)}, + data_sources={"daily": [{"name": "eastmoney", "enabled": True}]}, + validation={}, performance={}, + ) + with patch("sanguo_data.scheduler.fetch_daily", return_value=pd.DataFrame({ + "date": ["2026-01-01"], "open": [10.0], "high": [10.0], + "low": [10.0], "close": [10.0], "volume": [100]})) as m_fetch, \ + patch("sanguo_data.scheduler.write_daily") as m_write: + report = run_daily_update(cfg, symbols=["600000"]) + assert m_fetch.call_count == 0 # 已 done,跳过 + assert isinstance(report, UpdateReport) + assert report.skipped == 1 diff --git a/tests/data/test_spike_vnpy44.py b/tests/data/test_spike_vnpy44.py new file mode 100644 index 0000000..b8f3248 --- /dev/null +++ b/tests/data/test_spike_vnpy44.py @@ -0,0 +1,22 @@ +"""Spike: 验证 vnpy 4.4.0 数据库接口。无 NAS 数据时可 skip。""" +import datetime +import pytest +from vnpy.trader.database import get_database +from vnpy.trader.constant import Interval, Exchange + + +def test_vnpy44_database_interface(): + """验证 vnpy 4.4.0 数据库接口存在性和签名.""" + db = get_database() + assert hasattr(db, "load_bar_data") + assert hasattr(db, "save_bar_data") + + # 测试 load_bar_data 签名 + bars = db.load_bar_data( + symbol="600000", + exchange=Exchange.SSE, + interval=Interval.DAILY, + start=datetime.datetime(2026, 1, 1), + end=datetime.datetime(2026, 6, 30), + ) + assert isinstance(bars, list) diff --git a/tests/data/test_validator.py b/tests/data/test_validator.py new file mode 100644 index 0000000..a96145e --- /dev/null +++ b/tests/data/test_validator.py @@ -0,0 +1,7 @@ +from sanguo_data.validator import validate_daily + +def test_validate_daily_keeps_good_rows(good_daily_df): + assert len(validate_daily(good_daily_df)) == 2 + +def test_validate_daily_drops_zero_price(bad_daily_df): + assert len(validate_daily(bad_daily_df)) == 0 \ No newline at end of file