banbox/cbpred
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

update

Browse Source
This commit is contained in:
anyongjin
2026-10-07 21:23:21 +08:00
parent 79b02d1c7d
commit 9c8e9d523a
11 changed files with 557 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files

117
tests/test_latest_predict.py
View File

@@ -1,4 +1,7 @@
import json
import os
import subprocess
import sys
import tempfile
import unittest
from pathlib import Path
@@ -9,15 +12,129 @@ 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,

100
tests/test_pipeline.py
View File

@@ -2,13 +2,15 @@ from __future__ import annotations
import json
import subprocess
import sys
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch
from unittest.mock import MagicMock, patch
import numpy as np
import pandas as pd
import yaml
from cbpred.pipeline import (
CSV_COLUMNS,
@@ -25,6 +27,9 @@ from cbpred.pipeline import (
run_training,
Symbol,
run_banbot_download,
detect_db_type,
effective_banbot_config,
prepare_data,
)
@@ -44,6 +49,99 @@ def make_hour_rows(start_ms: int, hours: int, phase: float = 0.0):
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)