claude_dev
bdad396243
原 2 测试测"读 parquet"(旧实现),实现早已改读 vnpy DbBarData→一直 latent fail (KeyError→防崩溃后 RuntimeError)。套 test_datareader 的 vnpy mock 模式重写: 测 load_bar_data 返回 DataFrame + start/end 委托 vnpy 过滤。 注:read_index_daily 切 dbbardata 根治待定(000300 不在 dbbardata/000905 exchange 约定不匹配),见 memory read-index-daily-dbbardata-migration-pending。
186 lines
6.5 KiB
Python
186 lines
6.5 KiB
Python
"""Tests for index downloader and read_index_daily functionality."""
|
|
import pandas as pd
|
|
import os
|
|
from unittest.mock import patch, MagicMock
|
|
from datetime import date
|
|
from pathlib import Path
|
|
import pytest
|
|
|
|
from sanguo_data.config import DataConfig
|
|
|
|
|
|
def test_download_index_writes_parquet(tmp_path):
|
|
"""Test that download_index writes parquet files with correct structure."""
|
|
# Sample data that baostock would return
|
|
sample_data = [
|
|
["2024-01-02", "3495.0", "3505.0", "3490.0", "3500.0", "100000"],
|
|
["2024-01-03", "3505.0", "3515.0", "3500.0", "3510.0", "120000"],
|
|
["2024-01-04", "3515.0", "3525.0", "3510.0", "3520.0", "110000"],
|
|
]
|
|
|
|
# Create a simple baostock mock
|
|
class MockBaostock:
|
|
class MockResult:
|
|
def __init__(self, data):
|
|
self.error_code = "success"
|
|
self.error_msg = "success"
|
|
self.data = data
|
|
self.fields = ["date", "open", "high", "low", "close", "volume"]
|
|
self.row_index = 0
|
|
|
|
def next(self):
|
|
if self.row_index < len(self.data):
|
|
row = self.data[self.row_index]
|
|
self.row_index += 1
|
|
return True
|
|
return False
|
|
|
|
def get_row_data(self):
|
|
return self.data[self.row_index - 1]
|
|
|
|
def login(self):
|
|
return self.MockResult([])
|
|
|
|
def logout(self):
|
|
return self.MockResult([])
|
|
|
|
def query_history_k_data_plus(self, *args, **kwargs):
|
|
return self.MockResult(sample_data)
|
|
|
|
# Patch baostock module
|
|
import sys
|
|
sys.modules["baostock"] = MockBaostock()
|
|
|
|
try:
|
|
# Import after patching
|
|
from sanguo_data.index_downloader import download_index
|
|
|
|
# Download index data
|
|
download_index("sh000300", 2024, 2024, str(tmp_path))
|
|
|
|
finally:
|
|
# Clean up the mock
|
|
del sys.modules["baostock"]
|
|
|
|
# Verify parquet file was created
|
|
expected_file = tmp_path / "2024" / "sh000300_daily.parquet"
|
|
assert expected_file.exists(), f"Expected parquet file at {expected_file}"
|
|
|
|
# Verify parquet content
|
|
df_read = pd.read_parquet(expected_file)
|
|
assert len(df_read) == 3
|
|
assert "close" in df_read.columns
|
|
assert "date" in df_read.columns
|
|
assert df_read["close"].iloc[0] == 3500.0
|
|
|
|
|
|
def test_download_index_clears_proxy(tmp_path):
|
|
"""Test that download_index clears proxy environment variables."""
|
|
# Set proxy variables
|
|
os.environ["http_proxy"] = "http://evil:8080"
|
|
os.environ["https_proxy"] = "https://evil:8080"
|
|
|
|
sample_data = [["2024-01-02", "3495.0", "3505.0", "3490.0", "3500.0", "100000"]]
|
|
|
|
# Create a simple baostock mock
|
|
class MockBaostock:
|
|
class MockResult:
|
|
def __init__(self, data):
|
|
self.error_code = "success"
|
|
self.error_msg = "success"
|
|
self.data = data
|
|
self.fields = ["date", "open", "high", "low", "close", "volume"]
|
|
self.row_index = 0
|
|
|
|
def next(self):
|
|
if self.row_index < len(self.data):
|
|
row = self.data[self.row_index]
|
|
self.row_index += 1
|
|
return True
|
|
return False
|
|
|
|
def get_row_data(self):
|
|
return self.data[self.row_index - 1]
|
|
|
|
def login(self):
|
|
return self.MockResult([])
|
|
|
|
def logout(self):
|
|
return self.MockResult([])
|
|
|
|
def query_history_k_data_plus(self, *args, **kwargs):
|
|
return self.MockResult(sample_data)
|
|
|
|
# Patch baostock module
|
|
import sys
|
|
sys.modules["baostock"] = MockBaostock()
|
|
|
|
try:
|
|
# Import after patching
|
|
from sanguo_data.index_downloader import download_index
|
|
|
|
# Download index data
|
|
download_index("sh000300", 2024, 2024, str(tmp_path))
|
|
|
|
finally:
|
|
# Clean up the mock
|
|
del sys.modules["baostock"]
|
|
|
|
# Verify proxy variables were cleared
|
|
assert "http_proxy" not in os.environ
|
|
assert "https_proxy" not in os.environ
|
|
|
|
|
|
def test_read_index_daily_reads_from_vnpy_db(tmp_path):
|
|
"""read_index_daily 经 vnpy get_database.load_bar_data 读指数日线,返回 DataFrame。
|
|
|
|
注:read_index_daily 数据源是 vnpy DbBarData 表(非 parquet);D2 根治拟切 dbbardata,
|
|
届时 mock 数据源随之更新。当前测 vnpy 路径(与 read_db_daily 同 mock 模式)。
|
|
"""
|
|
from types import SimpleNamespace
|
|
from datetime import datetime
|
|
from sanguo_data.datareader import read_index_daily
|
|
|
|
cfg = DataConfig(
|
|
data_paths={"vnpy_db": str(tmp_path / "q.db")},
|
|
data_sources={}, validation={}, performance={},
|
|
)
|
|
# vnpy load_bar_data 返回类 BarData 对象(有 datetime/open_price/... 属性)
|
|
bars = [
|
|
SimpleNamespace(datetime=datetime(2024, 1, 2), open_price=3495.0,
|
|
high_price=3505.0, low_price=3490.0, close_price=3500.0, volume=100000),
|
|
SimpleNamespace(datetime=datetime(2024, 1, 3), open_price=3505.0,
|
|
high_price=3515.0, low_price=3500.0, close_price=3510.0, volume=120000),
|
|
SimpleNamespace(datetime=datetime(2024, 1, 4), open_price=3515.0,
|
|
high_price=3525.0, low_price=3510.0, close_price=3520.0, volume=110000),
|
|
]
|
|
mock_db = MagicMock()
|
|
mock_db.load_bar_data.return_value = bars
|
|
# read_index_daily 内 lazy import,patch 源模块属性(同 test_datareader)
|
|
with patch("vnpy.trader.database.get_database", return_value=mock_db):
|
|
result = read_index_daily("sh000300", date(2024, 1, 1), date(2024, 12, 31), cfg)
|
|
|
|
assert len(result) == 3
|
|
assert "close" in result.columns
|
|
assert result["close"].iloc[0] == 3500.0
|
|
|
|
|
|
def test_read_index_daily_passes_date_range_to_vnpy(tmp_path):
|
|
"""read_index_daily 把 start/end 传给 vnpy load_bar_data(日期过滤在 vnpy 层)。"""
|
|
from datetime import datetime
|
|
from sanguo_data.datareader import read_index_daily
|
|
|
|
cfg = DataConfig(
|
|
data_paths={"vnpy_db": str(tmp_path / "q.db")},
|
|
data_sources={}, validation={}, performance={},
|
|
)
|
|
mock_db = MagicMock()
|
|
mock_db.load_bar_data.return_value = []
|
|
with patch("vnpy.trader.database.get_database", return_value=mock_db):
|
|
read_index_daily("sh000300", date(2024, 1, 1), date(2024, 6, 30), cfg)
|
|
|
|
# start=date→combine 00:00:00; end=date→combine datetime.max.time()(23:59:59.999999)
|
|
call = mock_db.load_bar_data.call_args
|
|
assert call.kwargs["start"] == datetime(2024, 1, 1, 0, 0, 0)
|
|
assert call.kwargs["end"] == datetime(2024, 6, 30, 23, 59, 59, 999999)
|