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

5
.gitignore vendored
View File

@@ -2,6 +2,10 @@
__pycache__/
*.py[cod]
.pytest_cache/
config/banbot.local.yml
.cbpred-latest-predict.lock
go.work
go.work.sum
data/raw/
data/manifest.json
artifacts/
@@ -9,4 +13,3 @@ artifacts_*/
reports/
predictions/
data_*/

2
README.md
View File

@@ -8,7 +8,7 @@
## 快速开始
运行环境需要 Python 3、Go、可访问的 banbot 源码和 banbot 配置中的 QuestDB。默认路径是：
运行环境需要 Python 3、Go、可访问的 banbot 源码和数据库。脚本优先使用已有 TimescaleDB 配置，支持 `--banbot-config`、`--dsn` 和 `--db-type timescale`。Linux 的 banbot 路径通常是：
```text
/data/ban/banbot

109
cbpred/latest_predict.py
View File

@@ -2,19 +2,34 @@ from __future__ import annotations
import argparse
import csv
import fcntl
import errno
import json
import os
import socket
import subprocess
import tempfile
import time
from datetime import datetime, timezone
from pathlib import Path
from urllib.parse import urlparse
if os.name == "nt":
import msvcrt
else:
import fcntl
import joblib
import pandas as pd
import yaml
from psycopg.conninfo import conninfo_to_dict
from .pipeline import load_database_url, parse_time_ms, predict_from_csv, prepare_data, timeframe_step_ms
from .pipeline import (
effective_banbot_config,
load_database_url,
parse_time_ms,
predict_from_csv,
prepare_data,
timeframe_step_ms,
)
REPO_ROOT = Path(__file__).resolve().parents[1]
@@ -32,6 +47,22 @@ OUTPUT_COLUMNS = (
)
def _acquire_update_lock(lock) -> None:
if os.name != "nt":
fcntl.flock(lock, fcntl.LOCK_EX)
return
# Windows locks a byte range; every caller locks the first byte.
while True:
lock.seek(0)
try:
msvcrt.locking(lock.fileno(), msvcrt.LK_NBLCK, 1)
return
except OSError as exc:
if exc.errno not in {errno.EACCES, errno.EAGAIN, errno.EDEADLK}:
raise
time.sleep(0.1)
def _utc_label(timestamp_ms: int) -> str:
return datetime.fromtimestamp(timestamp_ms / 1000, timezone.utc).strftime(
"%Y-%m-%d %H:%M:%S UTC"
@@ -39,8 +70,35 @@ def _utc_label(timestamp_ms: int) -> str:
def _database_endpoint(dsn: str) -> tuple[str, int]:
parsed = urlparse(dsn)
return parsed.hostname or "127.0.0.1", parsed.port or 8812
parameters = conninfo_to_dict(dsn)
host = parameters.get("host") or parameters.get("hostaddr") or "127.0.0.1"
port = parameters.get("port") or "5432"
return host.split(",")[0], int(port.split(",")[0] or 5432)
def default_banbot_dir() -> Path:
configured = os.environ.get("CBPRED_BANBOT_DIR")
if configured:
return Path(configured)
candidates = (
REPO_ROOT.parent.parent / "ban/banbot",
REPO_ROOT.parent / "banbot",
Path("/data/ban/banbot"),
)
return next((path for path in candidates if path.is_dir()), candidates[-1])
def default_banbot_config(banbot_dir: Path) -> Path:
configured = os.environ.get("CBPRED_BANBOT_CONFIG")
if configured:
return Path(configured)
candidates = (
REPO_ROOT / "config/banbot.local.yml",
banbot_dir.parent / "data/config.timescale.yml",
banbot_dir / "biz/config.timescale.yml",
banbot_dir / "biz/config.yml",
)
return next((path for path in candidates if path.is_file()), candidates[-1])
def output_columns_for_timeframe(timeframe: str) -> tuple[str, ...]:
@@ -98,28 +156,40 @@ def _validate_freshness(
def ensure_database(
banbot_config: Path,
dsn: str | None,
db_type: str = "auto",
) -> None:
database_url = load_database_url(banbot_config, dsn)
if db_type == "auto" and dsn is None:
config = yaml.safe_load(banbot_config.read_text(encoding="utf-8")) or {}
db_type = (config.get("database") or {}).get("db_type") or "auto"
# PostgreSQL connections are validated by prepare_data via libpq, which also
# supports Unix sockets, multiple hosts and connection-service definitions.
if db_type not in {"auto", "questdb"}:
return
parameters = conninfo_to_dict(database_url)
if any("," in parameters.get(key, "") for key in ("host", "hostaddr", "port")):
return
host, port = _database_endpoint(database_url)
if db_type == "auto" and port != 8812:
return
try:
with socket.create_connection((host, port), timeout=1):
return
except OSError:
pass
if port != 8812:
return
print("[latest-predict] QuestDB is unavailable; letting banbot start it", flush=True)
start_ms = int(datetime.now(timezone.utc).timestamp() * 1000) - 8 * 60 * 60 * 1000
end_ms = int(datetime.now(timezone.utc).timestamp() * 1000)
with tempfile.TemporaryDirectory(prefix="cbpred-db-bootstrap-") as output_dir:
with tempfile.TemporaryDirectory(prefix="cbpred-db-bootstrap-") as output_dir, effective_banbot_config(
banbot_config, database_url, "questdb"
) as effective_config:
subprocess.run(
[
"go",
"run",
"./cmd/cbpred-export",
"--config",
str(banbot_config),
str(effective_config),
"--data-dir",
str(banbot_config.parent),
"--pairs",
@@ -222,11 +292,12 @@ def build_predictions(data_dir: Path, artifact: Path, output_dir: Path) -> dict[
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description="Update banbot K-lines and predict with the current model")
parser.add_argument("--banbot-dir", type=Path, default=Path("/data/ban/banbot"))
parser.add_argument("--banbot-dir", type=Path, default=default_banbot_dir())
parser.add_argument(
"--banbot-config",
type=Path,
default=Path("/data/ban/banbot/biz/config.yml"),
default=None,
help="banbot config (default: local config or existing TimescaleDB config)",
)
parser.add_argument("--data-dir", type=Path, default=None)
parser.add_argument("--artifact", type=Path, default=None)
@@ -241,8 +312,11 @@ def build_parser() -> argparse.ArgumentParser:
parser.add_argument("--min-age-days", type=int, default=730)
parser.add_argument("--chunk-size", type=int, default=100)
parser.add_argument("--end", default=None, help="UTC date/ISO timestamp or milliseconds")
parser.add_argument("--dsn", default=None)
parser.add_argument("--db-type", choices=["auto", "questdb", "timescale"], default="auto")
parser.add_argument("--dsn", default=os.environ.get("CBPRED_DSN"))
parser.add_argument(
"--db-type", choices=["auto", "questdb", "timescale"],
default=os.environ.get("CBPRED_DB_TYPE", "auto"),
)
parser.add_argument("--exchange", default="binance")
parser.add_argument("--timeframe", choices=["8h", "1d"], default="8h")
return parser
@@ -250,6 +324,9 @@ def build_parser() -> argparse.ArgumentParser:
def main() -> None:
args = build_parser().parse_args()
args.banbot_config = (
args.banbot_config or default_banbot_config(args.banbot_dir)
).resolve()
data_dir = (args.data_dir or default_data_dir(args.timeframe)).resolve()
artifact = (args.artifact or default_artifact(args.timeframe)).resolve()
output_dir = (
@@ -258,8 +335,8 @@ def main() -> None:
_validate_artifact_timeframe(artifact, args.timeframe)
lock_path = args.banbot_config.parent / ".cbpred-latest-predict.lock"
with lock_path.open("a+") as lock:
fcntl.flock(lock, fcntl.LOCK_EX)
ensure_database(args.banbot_config, args.dsn)
_acquire_update_lock(lock)
ensure_database(args.banbot_config, args.dsn, args.db_type)
manifest = prepare_data(
out_dir=data_dir,
banbot_dir=args.banbot_dir,

68
cbpred/pipeline.py
View File

@@ -8,10 +8,12 @@ import re
import shlex
import subprocess
import sys
import tempfile
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Iterable
from typing import Any, Iterable, Iterator
import joblib
import numpy as np
@@ -141,11 +143,59 @@ def detect_db_type(conn, requested: str = "auto") -> str:
return requested
with conn.cursor() as cur:
try:
cur.execute("SELECT 1 FROM exsymbol_q LIMIT 1")
return "questdb"
except Exception:
cur.execute("SELECT to_regclass('exsymbol'), to_regclass('exsymbol_q')")
timescale_table, questdb_table = cur.fetchone()
except Exception as exc:
conn.rollback()
# QuestDB speaks the PostgreSQL wire protocol but lacks to_regclass.
message = str(exc).lower()
if getattr(exc, "sqlstate", None) != "42883" and not (
"to_regclass" in message
and any(text in message for text in ("unknown function", "does not exist", "unsupported", "no signature"))
):
raise
else:
if timescale_table is not None:
return "timescale"
if questdb_table is not None:
return "questdb"
return "timescale"
for table, backend in (("exsymbol", "timescale"), ("exsymbol_q", "questdb")):
try:
cur.execute(f"SELECT 1 FROM {table} LIMIT 1")
return backend
except Exception as exc:
conn.rollback()
message = str(exc).lower()
if getattr(exc, "sqlstate", None) != "42P01" and not (
table in message and any(text in message for text in ("does not exist", "not found"))
):
raise
return "timescale"
@contextmanager
def effective_banbot_config(config_path: Path, dsn: str, db_type: str) -> Iterator[Path]:
"""Give Python and Go the same database without changing the original config."""
config_path = config_path.resolve()
config = yaml.safe_load(config_path.read_text(encoding="utf-8")) or {}
database = dict(config.get("database") or {})
database.update(url=dsn, db_type=db_type)
config["database"] = database
config["market_type"] = "linear"
temporary_path = None
try:
# Keep relative assets and --datadir anchored to the original directory.
with tempfile.NamedTemporaryFile(
mode="w", encoding="utf-8", suffix=".yml", prefix=".cbpred-db-",
dir=config_path.parent, delete=False,
) as handle:
temporary_path = Path(handle.name)
yaml.safe_dump(config, handle, allow_unicode=True, sort_keys=False)
yield temporary_path
finally:
if temporary_path is not None:
temporary_path.unlink(missing_ok=True)
def fetch_symbols(
@@ -250,7 +300,10 @@ def run_banbot_download(
str(end_ms),
]
try:
subprocess.run(command, cwd=banbot_dir, check=True, stderr=subprocess.PIPE, text=True)
subprocess.run(
command, cwd=banbot_dir, check=True,
stderr=subprocess.PIPE, text=True, encoding="utf-8",
)
except subprocess.CalledProcessError as exc:
invalid_symbols = set(INVALID_SYMBOL_RE.findall(exc.stderr or ""))
rejected = [item for item in chunk if item.symbol in invalid_symbols]
@@ -391,11 +444,12 @@ def prepare_data(
)
finally:
conn.close()
with effective_banbot_config(banbot_config, database_url, actual_db_type) as effective_config:
if not skip_download:
symbols, unavailable = run_banbot_download(
symbols,
banbot_dir,
banbot_config,
effective_config,
history_start_ms,
end_ms,
chunk_size,
@@ -405,7 +459,7 @@ def prepare_data(
run_banbot_export(
symbols,
banbot_dir,
banbot_config,
effective_config,
raw_dir,
history_start_ms,
end_ms,

62
cmd/cbpred-export/main.go
View File

@@ -6,6 +6,7 @@ import (
"fmt"
"os"
"path/filepath"
"sort"
"strconv"
"strings"
@@ -64,12 +65,21 @@ func main() {
}
fail(err)
}
_, rows, err := orm.GetSeries(exs, *timeframe, *startMS, *endMS, 0, false)
_, rows, err := orm.GetSeries(exs, "1h", *startMS, *endMS, 0, false)
if err != nil {
fail(err)
}
bars := make([]*orm.SeriesOHLCV, 0, len(rows))
for _, row := range rows {
bar, err := row.OHLCV(exs)
if err != nil {
fail(err)
}
bars = append(bars, bar)
}
bars = aggregateHourly(bars, *timeframe)
name := fmt.Sprintf("%d_%s.csv", exs.ID, strings.NewReplacer("/", "_", ":", "_").Replace(pair))
if err := writeCSV(filepath.Join(*outDir, name), rows, exs); err != nil {
if err := writeCSV(filepath.Join(*outDir, name), bars); err != nil {
fail(err)
}
}
@@ -79,7 +89,47 @@ func isUnavailablePair(err *errs.Error) bool {
return err != nil && err.Code == errs.CodeNoMarketForPair
}
func writeCSV(path string, rows []*orm.DataSeries, exs *orm.ExSymbol) error {
// aggregateHourly emits only complete UTC buckets with exactly one bar per hour.
func aggregateHourly(rows []*orm.SeriesOHLCV, timeframe string) []*orm.SeriesOHLCV {
const hourMS int64 = 60 * 60 * 1000
hours := 8
if timeframe == "1d" {
hours = 24
}
stepMS := int64(hours) * hourMS
ordered := append([]*orm.SeriesOHLCV(nil), rows...)
sort.Slice(ordered, func(i, j int) bool { return ordered[i].Time < ordered[j].Time })
var output []*orm.SeriesOHLCV
for first := 0; first < len(ordered); {
bucket := ordered[first].Time / stepMS * stepMS
end := first + 1
for end < len(ordered) && ordered[end].Time/stepMS*stepMS == bucket {
end++
}
complete := end-first == hours
for i := first; complete && i < end; i++ {
complete = ordered[i].Time == bucket+int64(i-first)*hourMS
}
if complete {
bar := *ordered[first]
bar.Time, bar.EndMS, bar.TimeFrame, bar.Closed = bucket, bucket+stepMS, timeframe, true
for _, row := range ordered[first+1 : end] {
bar.High = max(bar.High, row.High)
bar.Low = min(bar.Low, row.Low)
bar.Close = row.Close
bar.Volume += row.Volume
bar.Quote += row.Quote
bar.BuyVolume += row.BuyVolume
bar.TradeNum += row.TradeNum
}
output = append(output, &bar)
}
first = end
}
return output
}
func writeCSV(path string, rows []*orm.SeriesOHLCV) error {
file, err := os.Create(path)
if err != nil {
return err
@@ -89,11 +139,7 @@ func writeCSV(path string, rows []*orm.DataSeries, exs *orm.ExSymbol) error {
if err := writer.Write(columns); err != nil {
return err
}
for _, row := range rows {
bar, err := row.OHLCV(exs)
if err != nil {
return err
}
for _, bar := range rows {
values := []string{
strconv.FormatInt(bar.Time, 10),
strconv.FormatFloat(bar.Open, 'f', -1, 64),

95
cmd/cbpred-export/main_test.go
View File

@@ -1,8 +1,13 @@
package main
import (
"encoding/csv"
"os"
"path/filepath"
"reflect"
"testing"
"github.com/banbox/banbot/orm"
"github.com/banbox/banexg/errs"
)
@@ -14,3 +19,93 @@ func TestIsUnavailablePair(t *testing.T) {
t.Fatal("non-market errors must not be skipped")
}
}
const testHourMS int64 = 3600000
func hourlyBars(startHour, count int) []*orm.SeriesOHLCV {
rows := make([]*orm.SeriesOHLCV, 0, count)
for i := startHour; i < startHour+count; i++ {
rows = append(rows, &orm.SeriesOHLCV{
Time: int64(i) * testHourMS, EndMS: int64(i+1) * testHourMS,
Open: float64(i + 10), High: float64(i + 20), Low: float64(i), Close: float64(i + 12),
Volume: 2, Quote: 3, BuyVolume: 1, TradeNum: 4,
})
}
return rows
}
func TestAggregateHourlyDailyAndCSV(t *testing.T) {
// Reversed inputs still produce UTC-aligned daily history in time order.
rows := hourlyBars(0, 72)
for i, j := 0, len(rows)-1; i < j; i, j = i+1, j-1 {
rows[i], rows[j] = rows[j], rows[i]
}
bars := aggregateHourly(rows, "1d")
if len(bars) != 3 {
t.Fatalf("daily bars = %d, want 3", len(bars))
}
for i, bar := range bars {
if bar.Time != int64(i*24)*testHourMS || bar.EndMS != int64((i+1)*24)*testHourMS || bar.TimeFrame != "1d" || !bar.Closed {
t.Fatalf("unexpected daily metadata: %+v", bar)
}
}
path := filepath.Join(t.TempDir(), "daily.csv")
if err := writeCSV(path, bars); err != nil {
t.Fatal(err)
}
file, err := os.Open(path)
if err != nil {
t.Fatal(err)
}
defer file.Close()
contents, err := csv.NewReader(file).ReadAll()
if err != nil {
t.Fatal(err)
}
if !reflect.DeepEqual(contents[0], columns) {
t.Fatalf("unexpected CSV columns: %v", contents[0])
}
want := []string{"0", "10", "43", "0", "35", "48", "72", "24", "96"}
if !reflect.DeepEqual(contents[1], want) {
t.Fatalf("CSV first bar = %v, want %v", contents[1], want)
}
}
func TestAggregateHourlyEightHour(t *testing.T) {
bars := aggregateHourly(hourlyBars(0, 24), "8h")
if len(bars) != 3 {
t.Fatalf("eight-hour bars = %d, want 3", len(bars))
}
for i, bar := range bars {
if bar.Time != int64(i*8)*testHourMS || bar.Open != float64(i*8+10) || bar.Close != float64(i*8+19) || bar.Volume != 16 || bar.Quote != 24 || bar.BuyVolume != 8 || bar.TradeNum != 32 {
t.Fatalf("unexpected eight-hour bar: %+v", bar)
}
}
}
func TestAggregateHourlyRejectsIncompleteBuckets(t *testing.T) {
missing := hourlyBars(0, 48)
missing = append(missing[:12:12], missing[13:]...)
gapWithDuplicate := hourlyBars(0, 48)
gapWithDuplicate[12] = gapWithDuplicate[11]
duplicate := hourlyBars(0, 48)
duplicate = append(duplicate, duplicate[12])
for name, rows := range map[string][]*orm.SeriesOHLCV{
"missing hour": missing,
"gap with duplicate": gapWithDuplicate,
"extra duplicate": duplicate,
"partial start and end": hourlyBars(1, 70),
"misaligned hour": append(hourlyBars(0, 23), &orm.SeriesOHLCV{Time: 23*testHourMS + 1}),
} {
t.Run(name, func(t *testing.T) {
bars := aggregateHourly(rows, "1d")
want := 1
if name == "misaligned hour" {
want = 0
}
if len(bars) != want || (want == 1 && bars[0].Time != 24*testHourMS) {
t.Fatalf("unexpected complete buckets: %+v", bars)
}
})
}
}

19
docs/usage.md
View File

@@ -9,9 +9,11 @@
- Python 3 和本项目的虚拟环境；
- `pandas`、`numpy`、`scikit-learn`、`joblib` 等项目依赖；
- Go 和 `/data/ban/banbot` 源码；
- banbot 配置中的 QuestDB 数据库，以及对 Binance 数据的访问。
- banbot 配置中的 TimescaleDB（优先）或 QuestDB 数据库，以及对 Binance 数据的访问。
默认配置文件是 `/data/ban/banbot/biz/config.yml`。如果路径不同，可通过 `--banbot-config` 和 `--banbot-dir` 覆盖。也可以直接传入 `--dsn`。
一键脚本优先使用 `config/banbot.local.yml`，其次使用 banbot 同级 `data/config.timescale.yml` 或 `biz/config.timescale.yml` 中的 TimescaleDB 配置，最后使用 `biz/config.yml`。本机 banbot 目录会自动查找，也可通过 `--banbot-dir`、`--banbot-config`，或 `CBPRED_BANBOT_DIR`、`CBPRED_BANBOT_CONFIG` 环境变量指定。
`--dsn`（或 `CBPRED_DSN`）覆盖数据库连接；`--db-type timescale`（或 `CBPRED_DB_TYPE=timescale`）明确使用 TimescaleDB。自动检测优先选择 `exsymbol` 表对应的 TimescaleDB。Python 筛选、Go 下载和导出使用同一连接和数据库类型；连接失败会报错。只有 QuestDB 配置会尝试自动启动 QuestDB。
安装：
@@ -28,14 +30,23 @@ python3 -m venv .venv
脚本调用 `python -m cbpred.latest_predict`，按以下顺序执行：
1. 检查配置中的 QuestDB 是否可连接；QuestDB 默认不可用时，通过 banbot 的导出命令触发其启动流程。
1. 检查所选数据库是否可连接；TimescaleDB 不可用时直接报错，QuestDB 配置可通过 banbot 的导出命令触发其启动流程。
2. 从数据库筛选 Binance `linear` 市场、结算币为 USDT、上市至少 730 天且未退市的品种。
3. 调用 banbot 下载最近 36 个月的 1H K 线。
4. 调用 banbot 导出请求周期的完整 K 线；默认是 8H，也可以通过 `--timeframe 1d` 导出日线。
4. 调用 banbot 读取 1H 历史，并按 UTC 边界聚合请求周期的完整 K 线；默认是 8H，也可以通过 `--timeframe 1d` 导出日线。缺少小时数据的周期会被排除。
5. 写入 `data/manifest.json` 和 `data/raw/*.csv`。
6. 加载 `artifacts/v1/model.joblib`，对清单中的每个品种使用最新完整 8H K 线预测。
7. 按最新 8H USDT 成交额降序写入 `predictions/latest_v1_8h/predictions.csv`，并写入 `report.json`。
分别生成 8H 和日线预测：
```bash
./scripts/update_and_predict.sh --timeframe 8h --db-type timescale
./scripts/update_and_predict.sh --timeframe 1d --db-type timescale
```
两个 CSV 分别位于 `predictions/latest_v1_8h/predictions.csv` 和 `predictions/latest_v1_1d/predictions.csv`。其中 `price_level_reached_80pct` 是预测的 80% 命中位置，`price_level_cap_80pct` 是回调上限参考位置。
脚本的默认参数可以覆盖：
```bash

2
go.mod
View File

@@ -111,4 +111,4 @@ require (
modernc.org/sqlite v1.39.1 // indirect
)
replace github.com/banbox/banbot => /data/ban/banbot
replace github.com/banbox/banbot => ../../ban/banbot

3
scripts/update_and_predict.sh Executable file → Normal file
View File

@@ -3,6 +3,9 @@ set -euo pipefail
REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
PYTHON_BIN="$REPO_ROOT/.venv/bin/python"
if [[ ! -x "$PYTHON_BIN" && -x "$REPO_ROOT/.venv/Scripts/python.exe" ]]; then
PYTHON_BIN="$REPO_ROOT/.venv/Scripts/python.exe"
fi
if [[ ! -x "$PYTHON_BIN" ]]; then
printf 'missing virtualenv Python: %s\n' "$PYTHON_BIN" >&2

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)