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()