Add offline workflow and coverage tests

2026-06-16 17:37:16 +08:00
parent 8d908477e2
commit 31baa18ce5
16 changed files with 2104 additions and 9 deletions
@@ -0,0 +1,306 @@
+"""Offline downloader contract tests with mocked data providers."""
+
+from __future__ import annotations
+
+import numpy as np
+import pandas as pd
+
+import data.downloader as downloader
+import pipeline.data.downloader as pipeline_downloader
+from data.downloader import download_daily, download_daily_batch
+from pipeline.common.schema import DATA_COLUMNS
+from pipeline.data.downloader import download_universe
+
+
+class _FakeResult:
+    def __init__(self, rows, error_code="0", error_msg=""):
+        self.rows = rows
+        self.error_code = error_code
+        self.error_msg = error_msg
+        self._idx = -1
+
+    def next(self):
+        self._idx += 1
+        return self._idx < len(self.rows)
+
+    def get_row_data(self):
+        return self.rows[self._idx]
+
+
+def _daily_batch_row(
+    date: str = "2024-01-02",
+    open_: str = "10",
+    high: str = "11",
+    low: str = "9",
+    close: str = "10.5",
+    preclose: str = "10",
+    volume: str = "1000",
+    amount: str = "10500",
+) -> list[str]:
+    return [
+        date,
+        open_,
+        high,
+        low,
+        close,
+        preclose,
+        volume,
+        amount,
+        "1.23",
+        "5.0",
+        "1",
+        "0",
+        "8.0",
+        "1.1",
+        "2.2",
+        "3.3",
+    ]
+
+
+def test_download_daily_uses_baostock_before_akshare_in_auto(monkeypatch):
+    calls: list[str] = []
+    expected = pd.DataFrame({
+        "symbol": ["sh600000"],
+        "date": ["2024-01-02"],
+        "open": [10.0],
+        "high": [11.0],
+        "low": [9.0],
+        "close": [10.5],
+        "volume": [1000.0],
+        "amount": [10500.0],
+    })
+
+    def fake_baostock(symbol, start, end, adjust):
+        calls.append("baostock")
+        return expected
+
+    def fake_akshare(symbol, start, end, adjust):
+        calls.append("akshare")
+        raise AssertionError("akshare should not be called after baostock succeeds")
+
+    monkeypatch.setattr(downloader, "_download_baostock", fake_baostock)
+    monkeypatch.setattr(downloader, "_download_akshare", fake_akshare)
+
+    result = download_daily("sh600000", "2024-01-02", "2024-01-02", source="auto")
+
+    assert calls == ["baostock"]
+    assert result["date"].tolist() == [pd.Timestamp("2024-01-02")]
+    assert result["close"].tolist() == [10.5]
+
+
+def test_download_daily_falls_back_to_akshare_when_baostock_empty(monkeypatch):
+    calls: list[str] = []
+    fallback = pd.DataFrame({
+        "symbol": ["sz000001"],
+        "date": ["2024-01-02"],
+        "open": [20.0],
+        "high": [21.0],
+        "low": [19.0],
+        "close": [20.5],
+        "volume": [2000.0],
+        "amount": [41000.0],
+    })
+
+    monkeypatch.setattr(
+        downloader,
+        "_download_baostock",
+        lambda symbol, start, end, adjust: calls.append("baostock") or None,
+    )
+    monkeypatch.setattr(
+        downloader,
+        "_download_akshare",
+        lambda symbol, start, end, adjust: calls.append("akshare") or fallback,
+    )
+
+    result = download_daily("sz000001", "2024-01-02", "2024-01-02", source="auto")
+
+    assert calls == ["baostock", "akshare"]
+    assert result["symbol"].tolist() == ["sz000001"]
+    assert result["date"].tolist() == [pd.Timestamp("2024-01-02")]
+
+
+def test_download_daily_batch_maps_rich_schema_and_vwap(monkeypatch):
+    query_calls: list[dict] = []
+    login_count = 0
+    logout_count = 0
+
+    def fake_login():
+        nonlocal login_count
+        login_count += 1
+
+    def fake_logout():
+        nonlocal logout_count
+        logout_count += 1
+
+    def fake_query(**kwargs):
+        query_calls.append(kwargs)
+        rows = [
+            _daily_batch_row(volume="1000", amount="10500"),
+            _daily_batch_row(date="2024-01-03", volume="0", amount="0"),
+        ]
+        return _FakeResult(rows)
+
+    monkeypatch.setattr(downloader.bs, "login", fake_login)
+    monkeypatch.setattr(downloader.bs, "logout", fake_logout)
+    monkeypatch.setattr(downloader.bs, "query_history_k_data_plus", fake_query)
+
+    [(symbol, frame)] = list(
+        download_daily_batch(
+            ["sh600000"],
+            "2024-01-02",
+            "2024-01-03",
+            adjust="hfq",
+        )
+    )
+
+    assert symbol == "sh600000"
+    assert query_calls[0]["code"] == "sh.600000"
+    assert query_calls[0]["adjustflag"] == "1"
+    assert login_count == 1
+    assert logout_count == 1
+    assert frame is not None
+    assert frame.columns.tolist() == [
+        "symbol", "date", "open", "high", "low", "close", "preclose",
+        "volume", "amount", "vwap", "turn", "pctChg", "tradestatus", "isST",
+        "peTTM", "pbMRQ", "psTTM", "pcfNcfTTM",
+    ]
+    assert np.isclose(frame["vwap"].iloc[0], 10.5)
+    assert pd.isna(frame["vwap"].iloc[1])
+    assert pd.api.types.is_datetime64_any_dtype(frame["date"])
+    assert pd.api.types.is_numeric_dtype(frame["tradestatus"])
+
+
+def test_download_daily_batch_relogs_and_retries_session_loss(monkeypatch):
+    responses = [
+        _FakeResult([], error_code="10002007", error_msg="用户未登录"),
+        _FakeResult([_daily_batch_row()]),
+    ]
+    login_count = 0
+    logout_count = 0
+
+    def fake_login():
+        nonlocal login_count
+        login_count += 1
+
+    def fake_logout():
+        nonlocal logout_count
+        logout_count += 1
+
+    monkeypatch.setattr(downloader.bs, "login", fake_login)
+    monkeypatch.setattr(downloader.bs, "logout", fake_logout)
+    monkeypatch.setattr(
+        downloader.bs,
+        "query_history_k_data_plus",
+        lambda **kwargs: responses.pop(0),
+    )
+
+    [(symbol, frame)] = list(download_daily_batch(["sh600000"], "2024-01-02", "2024-01-02"))
+
+    assert symbol == "sh600000"
+    assert frame is not None
+    assert len(frame) == 1
+    assert login_count == 2
+    assert logout_count == 2
+
+
+def test_download_daily_batch_uses_akshare_fallback_when_enabled(monkeypatch):
+    fallback = pd.DataFrame({
+        "symbol": ["sh600000"],
+        "date": ["2024-01-02"],
+        "open": [10.0],
+        "high": [11.0],
+        "low": [9.0],
+        "close": [10.5],
+        "volume": [1000.0],
+        "amount": [10500.0],
+    })
+
+    monkeypatch.setattr(downloader.bs, "login", lambda: None)
+    monkeypatch.setattr(downloader.bs, "logout", lambda: None)
+    monkeypatch.setattr(
+        downloader.bs,
+        "query_history_k_data_plus",
+        lambda **kwargs: _FakeResult([], error_code="1", error_msg="no data"),
+    )
+    monkeypatch.setattr(
+        downloader,
+        "_download_akshare",
+        lambda symbol, start, end, adjust: fallback.copy(),
+    )
+
+    [(symbol, frame)] = list(
+        download_daily_batch(
+            ["sh600000"],
+            "2024-01-02",
+            "2024-01-02",
+            akshare_fallback=True,
+        )
+    )
+
+    assert symbol == "sh600000"
+    assert frame is not None
+    assert frame["date"].tolist() == [pd.Timestamp("2024-01-02")]
+    assert frame["close"].tolist() == [10.5]
+
+
+def test_download_universe_writes_daily_partitions_from_mock_batch(tmp_path, monkeypatch):
+    batch_frame = pd.DataFrame({
+        "symbol": ["sh600000", "sh600000"],
+        "date": pd.to_datetime(["2024-01-02", "2024-02-01"]),
+        "open": [10.0, 11.0],
+        "high": [11.0, 12.0],
+        "low": [9.0, 10.0],
+        "close": [10.5, 11.5],
+        "preclose": [10.0, 10.5],
+        "volume": [1000.0, 1200.0],
+        "amount": [10500.0, 13800.0],
+        "vwap": [10.5, 11.5],
+        "turn": [1.0, 1.1],
+        "pctChg": [5.0, 9.5],
+        "tradestatus": [1, 1],
+        "isST": [0, 0],
+        "peTTM": [8.0, 8.1],
+        "pbMRQ": [1.1, 1.2],
+        "psTTM": [2.1, 2.2],
+        "pcfNcfTTM": [3.1, 3.2],
+    })
+
+    monkeypatch.setattr(
+        pipeline_downloader,
+        "_resolve_universe",
+        lambda universe, max_symbols=0: pd.DataFrame({
+            "symbol_id": ["sh600000", "sz000001"],
+            "symbol_name": ["PF Bank", "Ping An Bank"],
+        }),
+    )
+
+    def fake_batch(symbols, start, end, adjust="qfq"):
+        assert symbols == ["sh600000", "sz000001"]
+        assert adjust == "qfq"
+        yield "sh600000", batch_frame
+        yield "sz000001", None
+
+    monkeypatch.setattr(pipeline_downloader, "download_daily_batch", fake_batch)
+
+    stats = download_universe(
+        universe="toy",
+        start_date="2024-01-02",
+        end_date="2024-02-01",
+        output_dir=str(tmp_path),
+        chunk_size=1,
+    )
+
+    dataset_path = tmp_path / "toy"
+    written = pd.read_parquet(dataset_path).sort_values(["date", "symbol_id"]).reset_index(drop=True)
+    assert stats == {
+        "dataset_path": str(dataset_path),
+        "n_symbols": 1,
+        "n_requested": 2,
+        "n_rows": 2,
+        "date_min": "2024-01-02",
+        "date_max": "2024-02-01",
+    }
+    assert (dataset_path / "month=2024-01").exists()
+    assert (dataset_path / "month=2024-02").exists()
+    assert written[DATA_COLUMNS].columns.tolist() == DATA_COLUMNS
+    assert written["symbol_name"].tolist() == ["PF Bank", "PF Bank"]