84 lines
2.9 KiB
Python
84 lines
2.9 KiB
Python
"""Test data adapter module - BarData to AlphaLab polars conversion."""
|
|
import sys
|
|
import os
|
|
_VNPY_SRC = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "vnpy_v4.4.0"))
|
|
sys.path.insert(0, _VNPY_SRC)
|
|
|
|
import polars as pl
|
|
from datetime import datetime
|
|
from vnpy.trader.object import BarData
|
|
from vnpy.trader.constant import Exchange, Interval
|
|
|
|
|
|
def _make_bar(symbol, dt, close):
|
|
"""Helper to create test BarData."""
|
|
return BarData(
|
|
symbol=symbol,
|
|
exchange=Exchange.SSE,
|
|
datetime=dt,
|
|
interval=Interval.DAILY,
|
|
open_price=close,
|
|
high_price=close,
|
|
low_price=close,
|
|
close_price=close,
|
|
volume=1000,
|
|
gateway_name="TEST"
|
|
)
|
|
|
|
|
|
def test_convert_bars_to_alpha_df_columns():
|
|
"""Test that convert_bars_to_alpha_df produces correct column structure."""
|
|
bars = [_make_bar("600000", datetime(2024, 1, i + 1), 10.0 + i) for i in range(5)]
|
|
from sanguo_factor.data_adapter import convert_bars_to_alpha_df
|
|
df = convert_bars_to_alpha_df(bars)
|
|
|
|
assert isinstance(df, pl.DataFrame)
|
|
# SPIKE CORRECTION: AlphaLab.save_bar_data stores columns as: open, high, low, close (not open_price)
|
|
for col in ["vt_symbol", "datetime", "open", "high", "low", "close", "volume", "turnover", "open_interest"]:
|
|
assert col in df.columns, f"Missing column: {col}"
|
|
assert df.height == 5
|
|
|
|
|
|
def test_convert_bars_to_alpha_df_values():
|
|
"""Test that convert_bars_to_alpha_df correctly converts BarData values."""
|
|
bars = [_make_bar("600000", datetime(2024, 1, i + 1), 10.0 + i) for i in range(5)]
|
|
from sanguo_factor.data_adapter import convert_bars_to_alpha_df
|
|
df = convert_bars_to_alpha_df(bars)
|
|
|
|
# Check vt_symbol format
|
|
assert df["vt_symbol"][0] == "600000.SSE"
|
|
|
|
# Check price values (SPIKE CORRECTION: use open/high/low/close column names)
|
|
assert df["open"][0] == 10.0
|
|
assert df["close"][4] == 14.0
|
|
|
|
# Check datetime
|
|
assert df["datetime"][0] == datetime(2024, 1, 1)
|
|
|
|
|
|
def test_convert_empty_bars():
|
|
"""Test that empty bar list returns empty DataFrame with correct schema."""
|
|
from sanguo_factor.data_adapter import convert_bars_to_alpha_df
|
|
df = convert_bars_to_alpha_df([])
|
|
|
|
assert df.height == 0
|
|
# Should still have schema defined
|
|
assert len(df.columns) == 9 # vt_symbol, datetime, open, high, low, close, volume, turnover, open_interest
|
|
|
|
|
|
def test_save_alpha_lab_data():
|
|
"""Test save_alpha_lab_data creates AlphaLab and saves data."""
|
|
import tempfile
|
|
from pathlib import Path
|
|
from sanguo_factor.data_adapter import save_alpha_lab_data
|
|
|
|
bars = [_make_bar("600000", datetime(2024, 1, i + 1), 10.0 + i) for i in range(5)]
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
lab_path = Path(tmpdir) / "alpha_lab"
|
|
result_path = save_alpha_lab_data(bars, str(lab_path))
|
|
|
|
# Should return the daily_path
|
|
assert result_path is not None
|
|
assert "daily" in str(result_path)
|