33 lines
1.1 KiB
Python
33 lines
1.1 KiB
Python
|
import numpy as np
|
||
|
|
import pandas as pd
|
||
|
|
|
||
|
|
from aitrader.features.indicators import compute_features
|
||
|
|
from aitrader.features.orderbook import summarize_orderbook
|
||
|
|
|
||
|
|
|
||
|
|
def _fake_ohlcv(n=250, seed=42):
|
||
|
|
rng = np.random.default_rng(seed)
|
||
|
|
close = 100 + np.cumsum(rng.normal(0, 1, n))
|
||
|
|
high = close + rng.uniform(0, 1, n)
|
||
|
|
low = close - rng.uniform(0, 1, n)
|
||
|
|
open_ = close + rng.normal(0, 0.5, n)
|
||
|
|
vol = rng.uniform(1, 10, n)
|
||
|
|
ts = pd.date_range("2025-01-01", periods=n, freq="15min", tz="UTC")
|
||
|
|
return pd.DataFrame({"ts": ts, "open": open_, "high": high, "low": low, "close": close, "volume": vol})
|
||
|
|
|
||
|
|
|
||
|
|
def test_compute_features_returns_keys():
|
||
|
|
df = _fake_ohlcv()
|
||
|
|
f = compute_features(df)
|
||
|
|
assert {"close", "rsi14", "macd", "ema20", "ema50", "ema200", "atr14"}.issubset(f.keys())
|
||
|
|
assert 0 <= f["rsi14"] <= 100
|
||
|
|
|
||
|
|
|
||
|
|
def test_orderbook_summary():
|
||
|
|
ob = {"bids": [[100.0, 1.0], [99.5, 2.0]], "asks": [[100.5, 1.5], [101.0, 1.0]]}
|
||
|
|
s = summarize_orderbook(ob)
|
||
|
|
assert s["best_bid"] == 100.0
|
||
|
|
assert s["best_ask"] == 100.5
|
||
|
|
assert s["spread_bps"] > 0
|
||
|
|
assert -1 <= s["depth_imbalance"] <= 1
|