import json import os import subprocess import sys import tempfile import unittest from pathlib import Path from unittest.mock import patch import pandas as pd from cbpred.latest_predict import ( OUTPUT_COLUMNS, REPO_ROOT, _acquire_update_lock, _database_endpoint, build_parser, build_predictions, default_artifact, default_banbot_config, default_banbot_dir, default_data_dir, default_output_dir, ensure_database, output_columns_for_timeframe, ) class LatestPredictTests(unittest.TestCase): def test_database_url_without_port_uses_postgresql_default(self): self.assertEqual(_database_endpoint("postgresql://localhost/qdb"), ("localhost", 5432)) def test_database_endpoint_supports_libpq_connection_formats(self): for dsn in ( "host=timescale.example port=5433 dbname=qdb", "postgresql:///qdb?host=timescale.example&port=5433", "host=timescale.example,backup.example port=5433,5434 dbname=qdb", ): self.assertEqual(_database_endpoint(dsn), ("timescale.example", 5433)) def test_banbot_defaults_support_checkout_at_drive_root(self): with patch.dict(os.environ, {}, clear=True), patch( "cbpred.latest_predict.REPO_ROOT", Path(Path.cwd().anchor) / "cbpred" ): self.assertIsInstance(default_banbot_dir(), Path) def test_default_config_prefers_timescale_and_respects_explicit_environment(self): with tempfile.TemporaryDirectory() as directory, patch.dict(os.environ, {}, clear=True): root = Path(directory) banbot_dir = root / "banbot" (banbot_dir / "biz").mkdir(parents=True) (root / "data").mkdir() legacy = banbot_dir / "biz/config.yml" legacy.write_text("database: {}\n") timescale = root / "data/config.timescale.yml" timescale.write_text("database: {db_type: timescale}\n") with patch("cbpred.latest_predict.REPO_ROOT", root / "cbpred"): self.assertEqual(default_banbot_config(banbot_dir), timescale) with patch.dict(os.environ, {"CBPRED_BANBOT_CONFIG": str(legacy)}): self.assertEqual(default_banbot_config(banbot_dir), legacy) def test_parser_reads_database_environment(self): with patch.dict(os.environ, {"CBPRED_DSN": "postgresql://localhost/qdb", "CBPRED_DB_TYPE": "timescale"}): args = build_parser().parse_args([]) self.assertEqual(args.dsn, "postgresql://localhost/qdb") self.assertEqual(args.db_type, "timescale") def test_unavailable_timescale_never_bootstraps_questdb(self): with patch("cbpred.latest_predict.socket.create_connection", side_effect=OSError), patch( "cbpred.latest_predict.subprocess.run" ) as bootstrap: ensure_database(Path("unused.yml"), "postgresql://localhost:8812/qdb", "timescale") bootstrap.assert_not_called() def test_configured_timescale_never_bootstraps_questdb_in_auto_mode(self): with tempfile.TemporaryDirectory() as directory: config = Path(directory) / "config.yml" config.write_text("database:\n db_type: timescale\n url: postgresql://localhost:8812/qdb\n") with patch("cbpred.latest_predict.socket.create_connection", side_effect=OSError), patch( "cbpred.latest_predict.subprocess.run" ) as bootstrap: ensure_database(config, None) bootstrap.assert_not_called() def test_postgresql_backend_aliases_never_bootstrap_questdb(self): with patch("cbpred.latest_predict.socket.create_connection", side_effect=OSError), patch( "cbpred.latest_predict.subprocess.run" ) as bootstrap: for backend in ("timescaledb", "postgres", "postgresql", "postgresql+timescale"): ensure_database(Path("unused.yml"), "postgresql://localhost:8812/qdb", backend) bootstrap.assert_not_called() def test_multihost_dsn_is_validated_by_libpq_without_bootstrap(self): with patch("cbpred.latest_predict.socket.create_connection") as probe, patch( "cbpred.latest_predict.subprocess.run" ) as bootstrap: for backend in ("auto", "timescale"): ensure_database( Path("unused.yml"), "host=primary,secondary port=8812,5433 dbname=qdb", backend, ) probe.assert_not_called() bootstrap.assert_not_called() def test_update_lock_serializes_processes_and_releases_on_close(self): with tempfile.TemporaryDirectory() as directory: lock_path = Path(directory) / "update.lock" script = ( "import sys\n" "from cbpred.latest_predict import _acquire_update_lock\n" "with open(sys.argv[1], 'a+') as lock:\n" " print('ready', flush=True)\n" " _acquire_update_lock(lock)\n" " print('acquired', flush=True)\n" ) with lock_path.open("a+") as lock: _acquire_update_lock(lock) process = subprocess.Popen( [sys.executable, "-c", script, str(lock_path)], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, ) try: self.assertEqual(process.stdout.readline().strip(), "ready") with self.assertRaises(subprocess.TimeoutExpired): process.communicate(timeout=0.3) lock.close() stdout, stderr = process.communicate(timeout=10) self.assertEqual(process.returncode, 0, stderr) self.assertEqual(stdout.strip(), "acquired") finally: if process.poll() is None: process.kill() process.communicate() def test_output_columns_are_stable(self): self.assertEqual( OUTPUT_COLUMNS, ( "rank", "symbol", "latest_bar_start_utc", "close", "quote_volume_usdt_8h", "atr_rank_200_pct", "drawdown_point_pct", "drawdown_cap_80pct_pct", "price_level_reached_80pct", "price_level_cap_80pct", ), ) def test_output_columns_follow_timeframe(self): self.assertEqual(output_columns_for_timeframe("1d")[4], "quote_volume_usdt_1d") def test_default_output_dir_follows_timeframe(self): self.assertEqual(default_output_dir("8h"), REPO_ROOT / "predictions/latest_v1_8h") self.assertEqual(default_output_dir("1d"), REPO_ROOT / "predictions/latest_v1_1d") def test_default_inputs_follow_timeframe(self): self.assertEqual(default_data_dir("8h"), REPO_ROOT / "data") self.assertEqual(default_data_dir("1d"), REPO_ROOT / "data_1d") self.assertEqual(default_artifact("8h"), REPO_ROOT / "artifacts/v1/model.joblib") self.assertEqual(default_artifact("1d"), REPO_ROOT / "artifacts_1d/best/model.joblib") def test_batch_rejects_stale_prediction_data(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) (root / "item.csv").write_text("timestamp,quote,close\n0,1,1\n") (root / "manifest.json").write_text( json.dumps( { "timeframe": "1d", "end_ms": 2 * 24 * 60 * 60 * 1000, "symbols": [{"symbol": "TRAIN/USDT:USDT", "path": "item.csv"}], } ) ) with patch( "cbpred.latest_predict.joblib.load", return_value={"timeframe": "1d", "asset_profiles": {"TRAIN/USDT:USDT": {}}}, ): with self.assertRaisesRegex(RuntimeError, "incomplete 1d data"): build_predictions(root, root / "model.joblib", root / "output") def test_batch_skips_symbols_without_training_profiles(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) data_dir = root / "data" output_dir = root / "output" data_dir.mkdir() for symbol, filename in ( ("TRAIN/USDT:USDT", "train.csv"), ("STALE/USDT:USDT", "stale.csv"), ): pd.DataFrame( [{"timestamp": 1, "quote": 10.0, "close": 100.0}] ).to_csv(data_dir / filename, index=False) (data_dir / "manifest.json").write_text( json.dumps( { "symbols": [ {"symbol": "TRAIN/USDT:USDT", "path": "train.csv"}, {"symbol": "STALE/USDT:USDT", "path": "stale.csv"}, ] } ) ) result = { "symbol": "TRAIN/USDT:USDT", "timestamp": 1, "close": 100.0, "volatility_rank_200": 0.5, "eligible_for_model": False, } with patch( "cbpred.latest_predict.joblib.load", return_value={"asset_profiles": {"TRAIN/USDT:USDT": {}}}, ), patch( "cbpred.latest_predict.predict_from_csv", return_value=result ) as predict: summary = build_predictions(data_dir, root / "model.joblib", output_dir) predict.assert_called_once_with( root / "model.joblib", data_dir / "train.csv", symbol="TRAIN/USDT:USDT" ) self.assertEqual(summary["symbols"], 1) self.assertEqual(summary["skipped_symbols"], ["STALE/USDT:USDT"]) if __name__ == "__main__": unittest.main()