Prediction and model training work
This commit is contained in:
@@ -1,8 +1,72 @@
|
||||
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"),
|
||||
|
||||
Reference in New Issue
Block a user