322 lines
13 KiB
Python
322 lines
13 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import yaml
|
|
|
|
from cbpred.pipeline import (
|
|
CSV_COLUMNS,
|
|
DAY_MS,
|
|
EIGHT_H_MS,
|
|
aggregate_ohlcv,
|
|
aggregate_8h,
|
|
build_feature_dataset,
|
|
compute_indicators,
|
|
fit_asset_scaler,
|
|
_vectorized_window_features,
|
|
_window_feature_names,
|
|
_window_features,
|
|
run_training,
|
|
Symbol,
|
|
run_banbot_download,
|
|
detect_db_type,
|
|
effective_banbot_config,
|
|
prepare_data,
|
|
)
|
|
|
|
|
|
def make_hour_rows(start_ms: int, hours: int, phase: float = 0.0):
|
|
rows = []
|
|
price = 100.0 + phase
|
|
for index in range(hours):
|
|
timestamp = start_ms + index * 60 * 60 * 1000
|
|
drift = 0.0015 * np.sin(index / 17 + phase)
|
|
open_ = price
|
|
close = price * (1 + drift)
|
|
high = max(open_, close) * 1.003
|
|
low = min(open_, close) * 0.997
|
|
rows.append((timestamp, open_, high, low, close, 1000 + index, 100000 + index * 10, 500, 10))
|
|
price = close
|
|
return rows
|
|
|
|
|
|
class PipelineTests(unittest.TestCase):
|
|
def test_download_reads_utf8_stderr_on_windows(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
downloader = root / "downloader.py"
|
|
downloader.write_text(
|
|
"import sys\n"
|
|
"sys.stderr.buffer.write('\\u4e00'.encode('utf-8'))\n",
|
|
encoding="utf-8",
|
|
)
|
|
symbol = Symbol(1, "binance", "linear", "BTC/USDT:USDT", 1, 0)
|
|
downloaded, unavailable = run_banbot_download(
|
|
[symbol], root, root / "config.yml", 1, 2,
|
|
banbot_cmd=f'"{sys.executable}" "{downloader}"',
|
|
)
|
|
self.assertEqual(downloaded, [symbol])
|
|
self.assertEqual(unavailable, [])
|
|
|
|
def test_auto_database_prefers_timescale_when_both_tables_exist(self):
|
|
conn = MagicMock()
|
|
cur = conn.cursor.return_value.__enter__.return_value
|
|
cur.fetchone.return_value = ("exsymbol", "exsymbol_q")
|
|
self.assertEqual(detect_db_type(conn), "timescale")
|
|
cur.execute.assert_called_once_with("SELECT to_regclass('exsymbol'), to_regclass('exsymbol_q')")
|
|
|
|
def test_auto_database_detects_questdb_postgres_tables(self):
|
|
conn = MagicMock()
|
|
conn.cursor.return_value.__enter__.return_value.fetchone.return_value = (None, "exsymbol_q")
|
|
self.assertEqual(detect_db_type(conn), "questdb")
|
|
|
|
def test_auto_database_falls_back_for_questdb_without_to_regclass(self):
|
|
conn = MagicMock()
|
|
cur = conn.cursor.return_value.__enter__.return_value
|
|
cur.execute.side_effect = [
|
|
RuntimeError("unknown function: to_regclass"),
|
|
RuntimeError("table does not exist [table=exsymbol]"),
|
|
None,
|
|
]
|
|
self.assertEqual(detect_db_type(conn), "questdb")
|
|
self.assertEqual(conn.rollback.call_count, 2)
|
|
|
|
def test_auto_database_does_not_hide_permission_errors(self):
|
|
conn = MagicMock()
|
|
conn.cursor.return_value.__enter__.return_value.execute.side_effect = RuntimeError("permission denied")
|
|
with self.assertRaisesRegex(RuntimeError, "permission denied"):
|
|
detect_db_type(conn)
|
|
|
|
def test_effective_config_is_removed_after_failure_and_original_unchanged(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
config_path = Path(directory) / "config.yml"
|
|
original = "database:\n url: original\n db_type: questdb\nassets: ./assets\n"
|
|
config_path.write_text(original)
|
|
with self.assertRaisesRegex(RuntimeError, "failed"):
|
|
with effective_banbot_config(config_path, "selected", "timescale") as effective:
|
|
self.assertEqual(effective.parent, config_path.parent)
|
|
self.assertTrue(effective.exists())
|
|
raise RuntimeError("failed")
|
|
self.assertFalse(effective.exists())
|
|
self.assertEqual(config_path.read_text(), original)
|
|
|
|
def test_prepare_propagates_selected_database_to_download_and_export(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
config_path = root / "config.yml"
|
|
original = "database:\n url: original\n db_type: questdb\n max_conn: 4\nmarket_type: spot\nassets: ./assets\n"
|
|
config_path.write_text(original)
|
|
symbol = Symbol(1, "binance", "linear", "BTC/USDT:USDT", 1, 0)
|
|
commands = []
|
|
config_paths = []
|
|
|
|
def run(command, **kwargs):
|
|
commands.append(command)
|
|
path = Path(command[command.index("--config") + 1])
|
|
config_paths.append(path)
|
|
effective = yaml.safe_load(path.read_text(encoding="utf-8"))
|
|
self.assertEqual(effective["database"], {"url": "selected-secret", "db_type": "timescale", "max_conn": 4})
|
|
self.assertEqual(effective["assets"], "./assets")
|
|
self.assertEqual(effective["market_type"], "linear")
|
|
self.assertEqual(path.parent, config_path.parent)
|
|
self.assertNotIn("selected-secret", command)
|
|
|
|
conn = MagicMock()
|
|
with patch("cbpred.pipeline.connect_database", return_value=conn) as connect, \
|
|
patch("cbpred.pipeline.detect_db_type", return_value="timescale"), \
|
|
patch("cbpred.pipeline.fetch_symbols", return_value=([symbol], [])), \
|
|
patch("cbpred.pipeline.subprocess.run", side_effect=run):
|
|
manifest = prepare_data(root / "data", root, config_path, dsn="selected-secret", end_ms=EIGHT_H_MS * 1000)
|
|
connect.assert_called_once_with("selected-secret")
|
|
self.assertEqual(len(commands), 2)
|
|
self.assertEqual(config_paths[0], config_paths[1])
|
|
self.assertFalse(config_paths[0].exists())
|
|
self.assertEqual(config_path.read_text(), original)
|
|
self.assertEqual(manifest["db_type"], "timescale")
|
|
|
|
def test_download_skips_invalid_current_symbol_and_retries_chunk(self):
|
|
valid = Symbol(1, "binance", "linear", "BTC/USDT:USDT", 1, 0)
|
|
stale = Symbol(2, "binance", "linear", "SXP/USDT:USDT", 1, 0)
|
|
calls = []
|
|
|
|
def run(command, **kwargs):
|
|
calls.append(command)
|
|
if len(calls) == 1:
|
|
raise subprocess.CalledProcessError(
|
|
1, command, stderr="Error: [InvalidSymbol] symbol SXP/USDT:USDT not found\n"
|
|
)
|
|
|
|
with patch("cbpred.pipeline.subprocess.run", side_effect=run):
|
|
downloaded, unavailable = run_banbot_download(
|
|
[valid, stale], Path("/banbot"), Path("/banbot/config.yml"), 1, 2
|
|
)
|
|
|
|
self.assertEqual(downloaded, [valid])
|
|
self.assertEqual(unavailable, [stale])
|
|
self.assertEqual(calls[0][calls[0].index("--pairs") + 1], "BTC/USDT:USDT,SXP/USDT:USDT")
|
|
self.assertEqual(calls[1][calls[1].index("--pairs") + 1], "BTC/USDT:USDT")
|
|
self.assertIn("market_type: linear", calls[0][calls[0].index("--config-data") + 1])
|
|
|
|
def test_aggregate_8h_drops_incomplete_buckets(self):
|
|
rows = make_hour_rows(0, 17)
|
|
frame = aggregate_8h(rows)
|
|
self.assertEqual(len(frame), 2)
|
|
self.assertEqual(frame.iloc[0]["timestamp"], 0)
|
|
self.assertAlmostEqual(frame.iloc[0]["quote"], sum(row[6] for row in rows[:8]))
|
|
|
|
def test_aggregate_1d_drops_incomplete_buckets(self):
|
|
rows = make_hour_rows(0, 49)
|
|
frame = aggregate_ohlcv(rows, "1d")
|
|
self.assertEqual(len(frame), 2)
|
|
self.assertEqual(frame.iloc[0]["timestamp"], 0)
|
|
self.assertEqual(frame.iloc[1]["timestamp"], DAY_MS)
|
|
|
|
def test_daily_feature_dataset_uses_daily_continuity(self):
|
|
raw = aggregate_ohlcv(make_hour_rows(0, 24 * 80), "1d")
|
|
profiles = {
|
|
"x": {
|
|
"asset_atr_price_top20_median": 0.01,
|
|
"asset_quote_volume_top20_median_log": 12.0,
|
|
}
|
|
}
|
|
scaler = fit_asset_scaler(profiles)
|
|
dataset, _, _ = build_feature_dataset(
|
|
{"x": raw},
|
|
1,
|
|
int(raw.iloc[55]["timestamp"]),
|
|
int(raw.iloc[-1]["timestamp"] + DAY_MS),
|
|
scaler,
|
|
volatility_rank_threshold=0.0,
|
|
step_ms=DAY_MS,
|
|
)
|
|
self.assertFalse(dataset.empty)
|
|
|
|
def test_future_target_changes_but_past_features_do_not(self):
|
|
start = 0
|
|
raw = aggregate_8h(make_hour_rows(start, 8 * 70))
|
|
indicators = compute_indicators(raw)
|
|
profiles = {"x": {"asset_atr_price_top20_median": 0.01, "asset_quote_volume_top20_median_log": 12.0}}
|
|
scaler = fit_asset_scaler(profiles)
|
|
one, names, _ = build_feature_dataset(
|
|
{"x": raw},
|
|
1,
|
|
int(raw.iloc[55]["timestamp"]),
|
|
int(raw.iloc[-1]["timestamp"] + EIGHT_H_MS),
|
|
scaler,
|
|
volatility_rank_threshold=0.0,
|
|
)
|
|
changed = raw.copy()
|
|
changed.loc[62:, "low"] *= 0.5
|
|
two, _, _ = build_feature_dataset(
|
|
{"x": changed},
|
|
1,
|
|
int(raw.iloc[55]["timestamp"]),
|
|
int(raw.iloc[-1]["timestamp"] + EIGHT_H_MS),
|
|
scaler,
|
|
volatility_rank_threshold=0.0,
|
|
)
|
|
anchor = int(raw.iloc[60]["timestamp"])
|
|
first = one[one["timestamp"] == anchor].iloc[0]
|
|
second = two[two["timestamp"] == anchor].iloc[0]
|
|
self.assertEqual(len(names), 18 * 10 + 2)
|
|
self.assertGreater(first["target_scale"], 0.0)
|
|
for name in names:
|
|
if name != "target":
|
|
self.assertEqual(first[name], second[name])
|
|
self.assertNotEqual(first["target"], second["target"])
|
|
|
|
rising = raw.copy()
|
|
anchor_close = float(rising.iloc[60]["close"])
|
|
for index in range(61, 67):
|
|
open_ = anchor_close * (1.01 + 0.002 * (index - 61))
|
|
close = open_ * 1.002
|
|
rising.loc[index, "open"] = open_
|
|
rising.loc[index, "low"] = open_ * 0.999
|
|
rising.loc[index, "high"] = close * 1.001
|
|
rising.loc[index, "close"] = close
|
|
rising_dataset, _, _ = build_feature_dataset(
|
|
{"x": rising},
|
|
1,
|
|
int(raw.iloc[55]["timestamp"]),
|
|
int(raw.iloc[-1]["timestamp"] + EIGHT_H_MS),
|
|
scaler,
|
|
volatility_rank_threshold=0.0,
|
|
)
|
|
rising_row = rising_dataset[rising_dataset["timestamp"] == anchor].iloc[0]
|
|
self.assertGreater(rising_row["target"], 0.0)
|
|
expected_target = (
|
|
np.quantile(rising.loc[61:66, "low"].to_numpy(dtype=float), 0.20)
|
|
/ anchor_close
|
|
- 1.0
|
|
) / rising_row["target_scale"]
|
|
self.assertAlmostEqual(rising_row["target"], expected_target)
|
|
|
|
def test_vectorized_features_match_window_features(self):
|
|
raw = aggregate_8h(make_hour_rows(0, 8 * 80))
|
|
indicators = compute_indicators(raw)
|
|
asset = {
|
|
"asset_atr_price_top20_median": 0.01,
|
|
"asset_quote_volume_top20_median_log": 12.0,
|
|
}
|
|
for stage in (1, 2, 3):
|
|
names = _window_feature_names(stage)
|
|
for anchor in (18, 30, 55):
|
|
window = indicators.iloc[anchor - 17 : anchor + 1]
|
|
expected = _window_features(window, stage, asset)
|
|
actual = _vectorized_window_features(
|
|
indicators, stage, np.array([anchor]), asset
|
|
)[0]
|
|
np.testing.assert_allclose(
|
|
actual,
|
|
np.array([expected[name] for name in names]),
|
|
rtol=1e-12,
|
|
atol=1e-12,
|
|
equal_nan=True,
|
|
)
|
|
|
|
def test_three_stage_training_writes_report_and_model(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
data_dir = root / "data"
|
|
raw_dir = data_dir / "raw"
|
|
raw_dir.mkdir(parents=True)
|
|
symbols = []
|
|
end_ms = 0
|
|
for idx in range(4):
|
|
symbol = f"COIN{idx}/USDT:USDT"
|
|
frame = aggregate_8h(make_hour_rows(0, 8 * 400, idx * 0.1))
|
|
path = raw_dir / f"{idx}.csv"
|
|
frame.to_csv(path, index=False)
|
|
end_ms = max(end_ms, int(frame.iloc[-1]["timestamp"] + EIGHT_H_MS))
|
|
symbols.append({"sid": idx, "symbol": symbol, "path": str(path.relative_to(data_dir)), "bars": len(frame)})
|
|
(data_dir / "manifest.json").write_text(
|
|
json.dumps({"end_ms": end_ms, "symbols": symbols})
|
|
)
|
|
report = run_training(
|
|
data_dir,
|
|
root / "artifacts",
|
|
end_ms=end_ms,
|
|
test_months=1,
|
|
n_estimators=4,
|
|
n_jobs=1,
|
|
)
|
|
self.assertIn(report["selected"]["stage"], {1, 2, 3})
|
|
self.assertEqual(len(report["stages"]), 3)
|
|
self.assertTrue((root / "artifacts" / "model.joblib").exists())
|
|
self.assertTrue((root / "artifacts" / "report.md").exists())
|
|
self.assertGreaterEqual(report["test"]["calibrated_q80_coverage"], 0.0)
|
|
self.assertEqual(report["protocol"]["volatility_filter"]["atr_rank_window"], 200)
|
|
self.assertAlmostEqual(report["protocol"]["volatility_filter"]["min_rank"], 0.0)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|