127 lines
5.9 KiB
Python
127 lines
5.9 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from cbpred.prediction_service import (
|
|
_prediction_command,
|
|
_ready_end_ms,
|
|
load_config,
|
|
run_ready_cycle,
|
|
send_email_notification,
|
|
send_http_notifications,
|
|
)
|
|
|
|
|
|
class PredictionServiceTests(unittest.TestCase):
|
|
def config(self, root: Path) -> dict:
|
|
return {
|
|
"_config_dir": root,
|
|
"scheduler": {"grace_seconds": 120, "state_file": "state.json"},
|
|
"prediction": {"working_dir": ".", "output_dir": "output"},
|
|
"notifications": {},
|
|
}
|
|
|
|
def test_load_config_expands_environment_and_requires_prediction(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
config_path = root / "service.yml"
|
|
config_path.write_text("prediction:\n dsn: '${CBPRED_TEST_DSN}'\n")
|
|
with patch.dict(os.environ, {"CBPRED_TEST_DSN": "postgres://test"}):
|
|
config = load_config(config_path)
|
|
self.assertEqual(config["prediction"]["dsn"], "postgres://test")
|
|
self.assertEqual(config["_config_dir"], root)
|
|
|
|
def test_command_uses_completed_candle_and_resolves_paths(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
config = self.config(root)
|
|
config["prediction"].update({"data_dir": "data", "history_months": 24})
|
|
command, cwd = _prediction_command(config, 28_800_000)
|
|
self.assertEqual(cwd, root)
|
|
self.assertEqual(command[:5], [os.sys.executable, "-m", "cbpred.latest_predict", "--end", "28800000"])
|
|
self.assertIn(str((root / "data").resolve()), command)
|
|
self.assertIn("24", command)
|
|
|
|
def test_cycle_runs_once_per_completed_candle_and_records_state(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
config = self.config(root)
|
|
csv_path = root / "output" / "predictions.csv"
|
|
csv_path.parent.mkdir()
|
|
csv_path.write_text("symbol\nBTC/USDT:USDT\n")
|
|
summary = {"symbols": 1, "eligible_predictions": 1, "csv": str(csv_path)}
|
|
completed_end = 28_800_000
|
|
with patch("cbpred.prediction_service.run_prediction", return_value=(summary, csv_path)) as run:
|
|
first = run_ready_cycle(config, now_ms=completed_end + 120_000)
|
|
second = run_ready_cycle(config, now_ms=completed_end + 121_000)
|
|
run.assert_called_once_with(config, completed_end)
|
|
self.assertIn("processed", first)
|
|
self.assertIn("already processed", second)
|
|
self.assertEqual(json.loads((root / "state.json").read_text())["last_successful_end_ms"], completed_end)
|
|
|
|
def test_cycle_waits_before_grace_period(self):
|
|
self.assertIsNone(_ready_end_ms(28_800_000 + 119_999, 120))
|
|
self.assertEqual(_ready_end_ms(28_800_000 + 120_000, 120), 28_800_000)
|
|
|
|
def test_failed_notification_does_not_advance_state(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
config = self.config(root)
|
|
csv_path = root / "output" / "predictions.csv"
|
|
csv_path.parent.mkdir()
|
|
csv_path.write_text("symbol\nBTC/USDT:USDT\n")
|
|
with patch(
|
|
"cbpred.prediction_service.run_prediction", return_value=({}, csv_path)
|
|
), patch(
|
|
"cbpred.prediction_service.send_http_notifications", side_effect=RuntimeError("unavailable")
|
|
):
|
|
with self.assertRaisesRegex(RuntimeError, "unavailable"):
|
|
run_ready_cycle(config, now_ms=28_920_000)
|
|
self.assertFalse((root / "state.json").exists())
|
|
|
|
def test_http_notification_posts_report_and_predictions(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
csv_path = root / "predictions.csv"
|
|
csv_path.write_text("symbol,rank\nBTC/USDT:USDT,1\n")
|
|
config = self.config(root)
|
|
config["notifications"] = {"http": [{"url": "https://example.test/hook", "headers": {"X-Key": "key"}}]}
|
|
response = MagicMock()
|
|
response.__enter__.return_value = response
|
|
with patch("cbpred.prediction_service.urlopen", return_value=response) as request:
|
|
count = send_http_notifications(config, {"symbols": 1}, csv_path, 28_800_000)
|
|
self.assertEqual(count, 1)
|
|
body = json.loads(request.call_args.args[0].data)
|
|
self.assertEqual(body["predictions"][0]["symbol"], "BTC/USDT:USDT")
|
|
self.assertEqual(request.call_args.args[0].headers["Content-type"], "application/json")
|
|
|
|
def test_email_notification_uses_ssl_client_and_attaches_csv(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
csv_path = root / "predictions.csv"
|
|
csv_path.write_text("symbol\nBTC/USDT:USDT\n")
|
|
config = self.config(root)
|
|
config["notifications"] = {
|
|
"email": {
|
|
"host": "smtp.example.test", "ssl": True, "username": "user", "password": "secret",
|
|
"from": "from@example.test", "to": ["to@example.test"],
|
|
}
|
|
}
|
|
client = MagicMock()
|
|
client.__enter__.return_value = client
|
|
with patch("cbpred.prediction_service.smtplib.SMTP_SSL", return_value=client) as smtp:
|
|
sent = send_email_notification(config, {"symbols": 1}, csv_path, 28_800_000)
|
|
self.assertTrue(sent)
|
|
smtp.assert_called_once()
|
|
client.login.assert_called_once_with("user", "secret")
|
|
client.send_message.assert_called_once()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|