from __future__ import annotations import json import subprocess import tempfile import unittest from pathlib import Path from unittest.mock import patch import numpy as np import pandas as pd 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, ) 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_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()