Yuxuan Yan
93 lines
2.8 KiB
Python
93 lines
2.8 KiB
Python
"""Offline tests for baostock-backed universe helpers."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pandas as pd
|
|
|
|
import data.universe as universe
|
|
|
|
|
|
class _FakeResult:
|
|
def __init__(self, rows, fields=None):
|
|
self.rows = rows
|
|
self.fields = fields or ["code", "name", "date"]
|
|
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 test_index_constituent_helpers_normalize_dotted_codes(monkeypatch):
|
|
calls: list[str] = []
|
|
|
|
monkeypatch.setattr(universe.bs, "login", lambda: calls.append("login"))
|
|
monkeypatch.setattr(universe.bs, "logout", lambda: calls.append("logout"))
|
|
monkeypatch.setattr(
|
|
universe.bs,
|
|
"query_hs300_stocks",
|
|
lambda: _FakeResult([
|
|
["sh.600000", "浦发银行", "2024-01-12"],
|
|
["sz.000001", "平安银行", "2024-01-12"],
|
|
]),
|
|
)
|
|
monkeypatch.setattr(
|
|
universe.bs,
|
|
"query_zz500_stocks",
|
|
lambda: _FakeResult([
|
|
["sh.600006", "东风汽车", "2024-01-12"],
|
|
]),
|
|
)
|
|
|
|
hs300 = universe.get_hs300_stocks()
|
|
zz500 = universe.get_zz500_stocks()
|
|
|
|
assert calls == ["login", "logout", "login", "logout"]
|
|
assert hs300["code"].tolist() == ["sh600000", "sz000001"]
|
|
assert zz500["code"].tolist() == ["sh600006"]
|
|
assert hs300["name"].tolist() == ["浦发银行", "平安银行"]
|
|
|
|
|
|
def test_get_all_stocks_walks_back_and_filters_to_listed_a_shares(monkeypatch):
|
|
fields = ["code", "tradeStatus", "code_name"]
|
|
responses = [
|
|
_FakeResult([], fields=fields),
|
|
_FakeResult(
|
|
[
|
|
["sh.600000", "1", "浦发银行"],
|
|
["sh.688001", "1", "华兴源创"],
|
|
["sz.000001", "1", "平安银行"],
|
|
["sz.300750", "1", "宁德时代"],
|
|
["sz.399001", "1", "深证成指"],
|
|
["sz.200001", "1", "深物业B"],
|
|
["bj.430047", "1", "北交所样本"],
|
|
],
|
|
fields=fields,
|
|
),
|
|
]
|
|
query_days: list[str] = []
|
|
|
|
monkeypatch.setattr(universe.bs, "login", lambda: None)
|
|
monkeypatch.setattr(universe.bs, "logout", lambda: None)
|
|
|
|
def fake_query_all_stock(day):
|
|
query_days.append(day)
|
|
return responses.pop(0)
|
|
|
|
monkeypatch.setattr(universe.bs, "query_all_stock", fake_query_all_stock)
|
|
|
|
result = universe.get_all_stocks("2024-01-07")
|
|
|
|
assert query_days == ["2024-01-07", "2024-01-06"]
|
|
assert result.columns.tolist() == ["code", "name"]
|
|
assert result["code"].tolist() == [
|
|
"sh600000",
|
|
"sh688001",
|
|
"sz000001",
|
|
"sz300750",
|
|
]
|
|
assert result["name"].tolist() == ["浦发银行", "华兴源创", "平安银行", "宁德时代"]
|