236 lines
9.9 KiB
Python
236 lines
9.9 KiB
Python
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()
|