80 lines
2.6 KiB
Python
80 lines
2.6 KiB
Python
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
|
|
from trading_bot.models.prediction import FEATURE_COLUMNS, predict_latest_probability
|
|
|
|
|
|
class FakeModel:
|
|
def predict_proba(self, model_input):
|
|
return np.array([[0.1, 0.9]])
|
|
|
|
|
|
def test_predict_latest_probability_refreshes_recent_market_data(
|
|
monkeypatch,
|
|
) -> None:
|
|
refresh_calls: dict[str, object] = {}
|
|
|
|
def fake_refresh_recent_data(
|
|
output_dir: Path,
|
|
symbols: dict[str, str] | None = None,
|
|
api_key: str | None = None,
|
|
secret_key: str | None = None,
|
|
) -> None:
|
|
refresh_calls["symbols"] = list((symbols or {}).keys())
|
|
refresh_calls["output_dir"] = output_dir
|
|
refresh_calls["api_key"] = api_key
|
|
refresh_calls["secret_key"] = secret_key
|
|
|
|
monkeypatch.setattr(
|
|
"trading_bot.models.prediction.refresh_recent_market_data",
|
|
fake_refresh_recent_data,
|
|
)
|
|
monkeypatch.setattr(
|
|
"trading_bot.models.prediction.load_model",
|
|
lambda model_path: FakeModel(),
|
|
)
|
|
monkeypatch.setattr(
|
|
"trading_bot.models.prediction.load_feature_metadata",
|
|
lambda metadata_path: {"feature_cols": FEATURE_COLUMNS, "p_base": 0.5},
|
|
)
|
|
monkeypatch.setattr(
|
|
"trading_bot.models.prediction.get_latest_inference_features",
|
|
lambda raw_data_dir, symbols=None, max_age_days=1: (
|
|
pd.Timestamp("2026-08-03"),
|
|
pd.DataFrame(
|
|
[[0.01, 0.02, 0.03, 0.04, 0.05, 0.06, 0.07, 0.08]],
|
|
columns=FEATURE_COLUMNS,
|
|
index=[pd.Timestamp("2026-08-03")],
|
|
),
|
|
),
|
|
)
|
|
|
|
result = predict_latest_probability(
|
|
model_path=Path("notebooks/models/spy_xgb_v1.json"),
|
|
metadata_path=Path("notebooks/models/spy_xgb_v1_meta.json"),
|
|
raw_data_dir=Path("data/alpaca/daily"),
|
|
api_key="test-key",
|
|
secret_key="test-secret",
|
|
fetch_recent_data=True,
|
|
)
|
|
|
|
assert refresh_calls["output_dir"] == Path("data/alpaca/daily")
|
|
assert refresh_calls["symbols"] == ["SPY", "VIXY", "TLT", "USO"]
|
|
assert refresh_calls["api_key"] == "test-key"
|
|
assert refresh_calls["secret_key"] == "test-secret"
|
|
assert result.probability == 0.9
|
|
|
|
|
|
def test_predict_latest_probability_returns_probability_between_zero_and_one() -> None:
|
|
result = predict_latest_probability(
|
|
model_path=Path("notebooks/models/spy_xgb_v1.json"),
|
|
metadata_path=Path("notebooks/models/spy_xgb_v1_meta.json"),
|
|
raw_data_dir=Path("data/alpaca/daily"),
|
|
)
|
|
|
|
assert set(result.feature_columns).issubset(FEATURE_COLUMNS)
|
|
assert 0.0 <= result.probability <= 1.0
|
|
assert result.probability > 0.0
|