Files
cbpred/tests/test_latest_predict.py
2026-08-29 11:00:27 +08:00

91 lines
3.0 KiB
Python

import json
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,
build_predictions,
default_output_dir,
output_columns_for_timeframe,
)
class LatestPredictTests(unittest.TestCase):
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_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()