Yuxuan Yan
528620b271
- pyproject.toml: fail_under 80 → 95 - test_alpha: +79 lines - test_cli_workflow: +226 lines - test_derived: +121 lines - test_downloader_contracts: +169 lines - test_features: +16 lines - test_minute_downloader: +81 lines - test_portfolio: +208 lines
377 lines
12 KiB
Python
377 lines
12 KiB
Python
"""Tests for daily derived-data ingestion and plugins."""
|
|
|
|
import textwrap
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pytest
|
|
from click.testing import CliRunner
|
|
|
|
from cli import cli
|
|
from pipeline.alpha.compute import join_feature_frames
|
|
from pipeline.derived.base import BaseDerivedData
|
|
from pipeline.derived.compute import compute_derived, read_derived_frame, validate_derived_frame
|
|
from pipeline.derived.registry import (
|
|
available_derived,
|
|
get_derived,
|
|
load_derived_module,
|
|
register_derived,
|
|
)
|
|
|
|
|
|
def _daily_bars() -> pd.DataFrame:
|
|
return pd.DataFrame({
|
|
"symbol_id": ["sh600000", "sz000001", "sh600000"],
|
|
"date": pd.to_datetime(["2024-01-02", "2024-01-02", "2024-01-03"]),
|
|
"open": [10.0, 20.0, 11.0],
|
|
"close": [10.5, 20.5, 11.5],
|
|
"volume": [1000.0, 2000.0, 1200.0],
|
|
})
|
|
|
|
|
|
def _minute_bars() -> pd.DataFrame:
|
|
return pd.DataFrame({
|
|
"symbol_id": ["sh600000", "sh600000", "sz000001"],
|
|
"datetime": pd.to_datetime([
|
|
"2024-01-02 09:35:00",
|
|
"2024-01-02 09:40:00",
|
|
"2024-01-02 09:35:00",
|
|
]),
|
|
"date": pd.to_datetime(["2024-01-02", "2024-01-02", "2024-01-02"]),
|
|
"time": ["09:35:00", "09:40:00", "09:35:00"],
|
|
"open": [10.0, 10.5, 20.0],
|
|
"high": [11.0, 12.0, 21.0],
|
|
"low": [9.0, 10.0, 19.0],
|
|
"close": [10.5, 11.0, 20.5],
|
|
"volume": [100.0, 300.0, 200.0],
|
|
"amount": [1000.0, 3300.0, 4100.0],
|
|
})
|
|
|
|
|
|
def test_validate_derived_frame_normalizes_and_sorts():
|
|
result = validate_derived_frame(pd.DataFrame({
|
|
"symbol_id": ["sz000001", "sh600000"],
|
|
"date": ["2024-01-02 15:00:00", "2024-01-02 09:30:00"],
|
|
"custom_value": [2.0, 1.0],
|
|
}))
|
|
|
|
assert result["symbol_id"].tolist() == ["sh600000", "sz000001"]
|
|
assert result["date"].tolist() == [
|
|
pd.Timestamp("2024-01-02"),
|
|
pd.Timestamp("2024-01-02"),
|
|
]
|
|
|
|
|
|
def test_validate_derived_frame_rejects_missing_keys():
|
|
with pytest.raises(ValueError, match="missing required"):
|
|
validate_derived_frame(pd.DataFrame({"symbol_id": ["sh600000"], "x": [1.0]}))
|
|
|
|
|
|
def test_validate_derived_frame_rejects_duplicate_normalized_keys():
|
|
with pytest.raises(ValueError, match="duplicate symbol_id,date"):
|
|
validate_derived_frame(pd.DataFrame({
|
|
"symbol_id": ["sh600000", "sh600000"],
|
|
"date": ["2024-01-02 09:30:00", "2024-01-02 15:00:00"],
|
|
"x": [1.0, 2.0],
|
|
}))
|
|
|
|
|
|
def test_validate_derived_frame_rejects_duplicate_columns():
|
|
bad = pd.DataFrame(
|
|
[["sh600000", pd.Timestamp("2024-01-02"), 1.0, 2.0]],
|
|
columns=["symbol_id", "date", "dup", "dup"],
|
|
)
|
|
with pytest.raises(ValueError, match="duplicate columns"):
|
|
validate_derived_frame(bad)
|
|
|
|
|
|
def test_validate_derived_frame_rejects_non_numeric_values():
|
|
with pytest.raises(ValueError, match="numeric"):
|
|
validate_derived_frame(pd.DataFrame({
|
|
"symbol_id": ["sh600000"],
|
|
"date": [pd.Timestamp("2024-01-02")],
|
|
"bad": ["not numeric"],
|
|
}))
|
|
|
|
|
|
def test_validate_derived_frame_rejects_missing_value_columns():
|
|
with pytest.raises(ValueError, match="at least one value column"):
|
|
validate_derived_frame(pd.DataFrame({
|
|
"symbol_id": ["sh600000"],
|
|
"date": [pd.Timestamp("2024-01-02")],
|
|
}))
|
|
|
|
|
|
def test_read_derived_frame_rejects_empty_csv(tmp_path):
|
|
empty_csv = tmp_path / "empty.csv"
|
|
empty_csv.write_text("")
|
|
|
|
with pytest.raises(ValueError, match="CSV input is empty"):
|
|
read_derived_frame(empty_csv)
|
|
|
|
|
|
def test_compute_derived_rejects_missing_inputs():
|
|
with pytest.raises(ValueError, match="requires --daily-path or --minute-path"):
|
|
compute_derived("minute_daily_summary")
|
|
|
|
|
|
def test_derived_ingest_cli_accepts_csv_and_parquet(tmp_path):
|
|
runner = CliRunner()
|
|
source = pd.DataFrame({
|
|
"symbol_id": ["sz000001", "sh600000"],
|
|
"date": ["2024-01-02", "2024-01-02"],
|
|
"custom_value": [2.0, 1.0],
|
|
})
|
|
csv_path = tmp_path / "custom.csv"
|
|
parquet_path = tmp_path / "custom.pq"
|
|
out_dir = tmp_path / "derived"
|
|
source.to_csv(csv_path, index=False)
|
|
source.to_parquet(parquet_path, index=False)
|
|
|
|
csv_result = runner.invoke(cli, [
|
|
"derived", "ingest",
|
|
"--input-path", str(csv_path),
|
|
"--derived-name", "csv_custom",
|
|
"--output-dir", str(out_dir),
|
|
])
|
|
assert csv_result.exit_code == 0, csv_result.output
|
|
|
|
parquet_result = runner.invoke(cli, [
|
|
"derived", "ingest",
|
|
"--input-path", str(parquet_path),
|
|
"--derived-name", "parquet_custom",
|
|
"--output-dir", str(out_dir),
|
|
])
|
|
assert parquet_result.exit_code == 0, parquet_result.output
|
|
|
|
written = pd.read_parquet(out_dir / "csv_custom.pq")
|
|
assert written["symbol_id"].tolist() == ["sh600000", "sz000001"]
|
|
assert (out_dir / "parquet_custom.pq").exists()
|
|
|
|
|
|
def test_derived_validate_cli_rejects_duplicate_csv_columns(tmp_path):
|
|
runner = CliRunner()
|
|
csv_path = tmp_path / "bad.csv"
|
|
csv_path.write_text("symbol_id,date,x,x\nsh600000,2024-01-02,1.0,2.0\n")
|
|
|
|
result = runner.invoke(cli, [
|
|
"derived", "validate",
|
|
"--input-path", str(csv_path),
|
|
])
|
|
|
|
assert result.exit_code != 0
|
|
assert "duplicate columns" in result.output
|
|
|
|
|
|
def test_derived_ingest_cli_wraps_validation_errors(tmp_path):
|
|
runner = CliRunner()
|
|
csv_path = tmp_path / "bad_ingest.csv"
|
|
csv_path.write_text("symbol_id,date\nsh600000,2024-01-02\n")
|
|
|
|
result = runner.invoke(cli, [
|
|
"derived",
|
|
"ingest",
|
|
"--input-path",
|
|
str(csv_path),
|
|
"--derived-name",
|
|
"bad",
|
|
"--output-dir",
|
|
str(tmp_path / "derived"),
|
|
])
|
|
|
|
assert result.exit_code != 0
|
|
assert "at least one value column" in result.output
|
|
|
|
|
|
def test_external_derived_plugin_loads_filters_params_and_uses_inputs(tmp_path):
|
|
module_path = tmp_path / "external_derived.py"
|
|
module_path.write_text(textwrap.dedent('''
|
|
import pandas as pd
|
|
from pipeline.derived.base import BaseDerivedData
|
|
from pipeline.derived.registry import register_derived
|
|
|
|
@register_derived
|
|
class FlexibleDerived(BaseDerivedData):
|
|
name = "flexible_derived_test"
|
|
|
|
def __init__(self, scale: float = 1.0):
|
|
self.scale = scale
|
|
|
|
def compute(self, daily=None, minute=None) -> pd.DataFrame:
|
|
result = None
|
|
if daily is not None:
|
|
result = daily[["symbol_id", "date", "close"]].copy()
|
|
result["daily_scaled_close"] = result.pop("close") * self.scale
|
|
if minute is not None:
|
|
minute_out = (
|
|
minute.groupby(["symbol_id", "date"], as_index=False)["volume"]
|
|
.sum()
|
|
.rename(columns={"volume": "minute_volume_sum"})
|
|
)
|
|
minute_out["minute_volume_sum"] *= self.scale
|
|
result = minute_out if result is None else result.merge(
|
|
minute_out, on=["symbol_id", "date"], how="left"
|
|
)
|
|
return result
|
|
'''))
|
|
|
|
load_derived_module(str(module_path))
|
|
assert "flexible_derived_test" in available_derived()
|
|
|
|
instance = get_derived("flexible_derived_test", scale=2.0, ignored=99)
|
|
assert instance.scale == 2.0
|
|
assert not hasattr(instance, "ignored")
|
|
|
|
daily_result = compute_derived(
|
|
"flexible_derived_test",
|
|
daily=_daily_bars(),
|
|
scale=2.0,
|
|
ignored=99,
|
|
)
|
|
assert "daily_scaled_close" in daily_result.columns
|
|
assert np.isclose(daily_result["daily_scaled_close"].iloc[0], 21.0)
|
|
|
|
minute_result = compute_derived(
|
|
"flexible_derived_test",
|
|
minute=_minute_bars(),
|
|
scale=2.0,
|
|
)
|
|
assert "minute_volume_sum" in minute_result.columns
|
|
assert np.isclose(
|
|
minute_result.loc[minute_result["symbol_id"] == "sh600000", "minute_volume_sum"].iloc[0],
|
|
800.0,
|
|
)
|
|
|
|
both_result = compute_derived(
|
|
"flexible_derived_test",
|
|
daily=_daily_bars(),
|
|
minute=_minute_bars(),
|
|
scale=1.0,
|
|
)
|
|
assert {"daily_scaled_close", "minute_volume_sum"}.issubset(both_result.columns)
|
|
|
|
|
|
def test_derived_registry_rejects_bad_plugins_and_load_failures(tmp_path, monkeypatch):
|
|
with pytest.raises(TypeError):
|
|
register_derived(object) # type: ignore[arg-type]
|
|
|
|
with pytest.raises(ValueError, match="non-empty"):
|
|
@register_derived
|
|
class NoNameDerived(BaseDerivedData):
|
|
def compute(self, daily=None, minute=None):
|
|
return pd.DataFrame()
|
|
|
|
@register_derived
|
|
class _CoverageDerived(BaseDerivedData):
|
|
name = "_coverage_derived_registry"
|
|
|
|
def compute(self, daily=None, minute=None):
|
|
return pd.DataFrame({
|
|
"symbol_id": ["sh600000"],
|
|
"date": [pd.Timestamp("2024-01-02")],
|
|
"x": [1.0],
|
|
})
|
|
|
|
with pytest.raises(ValueError, match="already registered"):
|
|
@register_derived
|
|
class _CoverageDerivedDuplicate(BaseDerivedData):
|
|
name = "_coverage_derived_registry"
|
|
|
|
def compute(self, daily=None, minute=None):
|
|
return pd.DataFrame()
|
|
|
|
with pytest.raises(KeyError, match="Unknown derived data"):
|
|
get_derived("does_not_exist")
|
|
|
|
missing_path = tmp_path / "missing_derived.py"
|
|
with pytest.raises(FileNotFoundError):
|
|
load_derived_module(str(missing_path))
|
|
|
|
bad_path = tmp_path / "bad_derived.py"
|
|
bad_path.write_text("x = 1\n")
|
|
monkeypatch.setattr(
|
|
"pipeline.derived.registry.importlib.util.spec_from_file_location",
|
|
lambda *args, **kwargs: None,
|
|
)
|
|
with pytest.raises(ImportError, match="Cannot load derived data module"):
|
|
load_derived_module(str(bad_path))
|
|
|
|
load_derived_module("math")
|
|
instance = _CoverageDerived()
|
|
instance.scale = 2
|
|
assert repr(instance) == "_CoverageDerived(scale=2)"
|
|
|
|
|
|
def test_derived_compute_cli_writes_builtin_minute_summary(tmp_path):
|
|
runner = CliRunner()
|
|
minute_path = tmp_path / "minute.pq"
|
|
out_dir = tmp_path / "derived"
|
|
_minute_bars().to_parquet(minute_path, index=False)
|
|
|
|
result = runner.invoke(cli, [
|
|
"derived", "compute",
|
|
"--minute-path", str(minute_path),
|
|
"--derived-type", "minute_daily_summary",
|
|
"--derived-name", "minute_summary",
|
|
"--output-dir", str(out_dir),
|
|
])
|
|
|
|
assert result.exit_code == 0, result.output
|
|
written = pd.read_parquet(out_dir / "minute_summary.pq")
|
|
assert "minute_vwap" in written.columns
|
|
|
|
|
|
def test_minute_summary_uses_time_sort_and_daily_without_close():
|
|
minute = _minute_bars().drop(columns=["datetime"]).sample(frac=1.0, random_state=0)
|
|
daily = _daily_bars()[["symbol_id", "date", "open"]]
|
|
|
|
result = compute_derived(
|
|
"minute_daily_summary",
|
|
daily=daily,
|
|
minute=minute,
|
|
)
|
|
|
|
by_symbol = result.set_index(["symbol_id", "date"])
|
|
assert np.isclose(
|
|
by_symbol.loc[("sh600000", pd.Timestamp("2024-01-02")), "minute_intraday_return"],
|
|
0.10,
|
|
)
|
|
assert "minute_vwap_deviation" in result.columns
|
|
|
|
|
|
def test_alpha_feature_join_rejects_derived_column_collisions():
|
|
data = _daily_bars()
|
|
derived_a = data[["symbol_id", "date"]].copy()
|
|
derived_a["custom_value"] = 1.0
|
|
derived_b = data[["symbol_id", "date"]].copy()
|
|
derived_b["custom_value"] = 2.0
|
|
|
|
with pytest.raises(ValueError, match="conflict"):
|
|
join_feature_frames(data, [derived_a, derived_b])
|
|
|
|
close_collision = data[["symbol_id", "date"]].copy()
|
|
close_collision["close"] = 1.0
|
|
with pytest.raises(ValueError, match="conflict"):
|
|
join_feature_frames(data, [close_collision])
|
|
|
|
|
|
def test_legacy_feature_cli_delegates_to_derived_registry(tmp_path):
|
|
runner = CliRunner()
|
|
minute_path = tmp_path / "minute.pq"
|
|
out_dir = tmp_path / "features"
|
|
_minute_bars().to_parquet(minute_path, index=False)
|
|
|
|
list_result = runner.invoke(cli, ["feature", "list"])
|
|
assert list_result.exit_code == 0, list_result.output
|
|
assert "minute_daily_summary" in list_result.output
|
|
|
|
compute_result = runner.invoke(cli, [
|
|
"feature", "compute",
|
|
"--minute-path", str(minute_path),
|
|
"--feature-type", "minute_daily_summary",
|
|
"--feature-name", "minute_summary",
|
|
"--output-dir", str(out_dir),
|
|
])
|
|
assert compute_result.exit_code == 0, compute_result.output
|
|
assert (out_dir / "minute_summary.pq").exists()
|