"""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)