TradeBot/tests/test_time_series.py

from __future__ import annotations

import json

from crypto_spot_bot.models import Candle
from crypto_spot_bot.time_series import TimeSeriesForecaster


def _candles_from_returns(returns: list[float]) -> list[Candle]:
    close = 100.0
    candles = [
        Candle(
            timestamp=0,
            open=close,
            high=close * 1.001,
            low=close * 0.999,
            close=close,
            volume=100,
        )
    ]
    for index, ret in enumerate(returns, start=1):
        previous = close
        close = close * (2.718281828459045 ** ret)
        candles.append(
            Candle(
                timestamp=index,
                open=previous,
                high=max(previous, close) * 1.001,
                low=min(previous, close) * 0.999,
                close=close,
                volume=100,
            )
        )
    return candles


def _write_torch_gru_artifact(
    path,
    *,
    head_bias: float,
    validation_mae_percent: float = 0.02,
    baseline_mae_percent: float = 0.08,
    skill: float = 0.2,
) -> None:
    hidden_size = 2
    path.write_text(
        json.dumps(
            {
                "version": 2,
                "type": "pytorch_recurrent_forecaster",
                "symbols": {
                    "BTCUSDT": {
                        "model": "torch_gru",
                        "architecture": "gru",
                        "lookback": 8,
                        "hidden_size": hidden_size,
                        "num_layers": 1,
                        "mean": 0.0,
                        "scale": 0.001,
                        "clip": 8.0,
                        "validation_mae_percent": validation_mae_percent,
                        "baseline_mae_percent": baseline_mae_percent,
                        "skill": skill,
                        "state_dict": {
                            "weight_ih_l0": [[0.0] for _ in range(3 * hidden_size)],
                            "weight_hh_l0": [[0.0, 0.0] for _ in range(3 * hidden_size)],
                            "bias_ih_l0": [0.0 for _ in range(3 * hidden_size)],
                            "bias_hh_l0": [0.0 for _ in range(3 * hidden_size)],
                        },
                        "head_weight": [0.0, 0.0],
                        "head_bias": head_bias,
                    },
                },
            }
        ),
        encoding="utf-8",
    )


def test_time_series_forecaster_requires_torch_artifact(make_settings, tmp_path) -> None:
    settings = make_settings(
        tmp_path,
        time_series_min_candles=80,
        time_series_forecast_horizon=3,
    )
    returns = []
    value = 0.0003
    for _ in range(140):
        value = 0.00025 + value * 0.55
        returns.append(value)

    forecast = TimeSeriesForecaster(settings).forecast(_candles_from_returns(returns), symbol="BTCUSDT")

    assert forecast.usable is False
    assert forecast.model == "none"
    assert forecast.candidates == []
    assert "PyTorch" in forecast.reason


def test_time_series_forecaster_blocks_negative_torch_edge(make_settings, tmp_path) -> None:
    artifact_path = tmp_path / "lstm_forecaster.json"
    _write_torch_gru_artifact(artifact_path, head_bias=-0.8, validation_mae_percent=0.01)
    settings = make_settings(
        tmp_path,
        time_series_min_candles=80,
        time_series_forecast_horizon=3,
        time_series_min_edge_percent=0.03,
        time_series_lstm_model_path=artifact_path,
    )
    returns = [0.00015 if index % 4 else -0.00005 for index in range(140)]

    forecast = TimeSeriesForecaster(settings).forecast(_candles_from_returns(returns), symbol="BTCUSDT")

    assert forecast.usable is True
    assert forecast.model == "torch_gru"
    assert forecast.expected_return_percent < 0
    assert forecast.probability_up < 0.45
    assert forecast.block_entry is True


def test_time_series_forecaster_ignores_legacy_lstm_artifact(make_settings, tmp_path) -> None:
    artifact_path = tmp_path / "lstm_forecaster.json"
    artifact_path.write_text(
        json.dumps(
            {
                "version": 1,
                "type": "lstm_reservoir_ridge_params",
                "symbols": {
                    "BTCUSDT": {"lookback": 10, "units": 3, "ridge": 0.01},
                },
            }
        ),
        encoding="utf-8",
    )
    settings = make_settings(
        tmp_path,
        time_series_min_candles=80,
        time_series_lstm_enabled=True,
        time_series_lstm_model_path=artifact_path,
    )
    returns = [0.00012 if index % 3 else -0.00008 for index in range(140)]

    forecast = TimeSeriesForecaster(settings).forecast(_candles_from_returns(returns), symbol="BTCUSDT")

    assert forecast.usable is False
    assert forecast.model == "none"
    assert forecast.candidates == []
    assert "PyTorch" in forecast.reason


def test_time_series_forecaster_reads_torch_gru_artifact(make_settings, tmp_path) -> None:
    artifact_path = tmp_path / "lstm_forecaster.json"
    _write_torch_gru_artifact(artifact_path, head_bias=0.2)
    settings = make_settings(
        tmp_path,
        time_series_min_candles=80,
        time_series_lstm_enabled=True,
        time_series_lstm_model_path=artifact_path,
    )
    returns = [0.00015 if index % 4 else -0.00005 for index in range(140)]

    forecast = TimeSeriesForecaster(settings).forecast(_candles_from_returns(returns), symbol="BTCUSDT")

    assert forecast.usable is True
    assert forecast.model == "torch_gru"
    assert forecast.candidates == [{"model": "torch_gru", "mae_percent": 0.02}]
    assert forecast.expected_return_percent > 0
    assert forecast.probability_up > 0.5