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 test_time_series_forecaster_selects_positive_predictive_model(make_settings, tmp_path) -> None:
    settings = make_settings(
        tmp_path,
        time_series_min_candles=80,
        time_series_validation_window=24,
        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))

    assert forecast.usable is True
    assert forecast.model != "naive"
    assert forecast.expected_return_percent > 0
    assert forecast.probability_up > 0.5


def test_time_series_forecaster_blocks_negative_edge(make_settings, tmp_path) -> None:
    settings = make_settings(
        tmp_path,
        time_series_min_candles=80,
        time_series_validation_window=24,
        time_series_forecast_horizon=3,
        time_series_min_edge_percent=0.03,
    )
    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))

    assert forecast.usable is True
    assert forecast.expected_return_percent < 0
    assert forecast.block_entry is True


def test_time_series_forecaster_includes_lstm_candidate(make_settings, tmp_path) -> None:
    settings = make_settings(
        tmp_path,
        time_series_min_candles=80,
        time_series_validation_window=20,
        time_series_lstm_enabled=True,
        time_series_lstm_lookback=12,
        time_series_lstm_units=4,
    )
    returns = []
    for index in range(140):
        seasonal = 0.00018 if index % 5 in {0, 1, 2} else -0.00011
        returns.append(seasonal + 0.00002 * ((index % 7) - 3))

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

    assert forecast.usable is True
    assert any(candidate["model"] == "lstm" for candidate in forecast.candidates)


def test_time_series_forecaster_reads_lstm_artifact(make_settings, tmp_path) -> None:
    artifact_path = tmp_path / "lstm_forecaster.json"
    artifact_path.write_text(
        json.dumps(
            {
                "version": 1,
                "symbols": {
                    "BTCUSDT": {"lookback": 10, "units": 3, "ridge": 0.01},
                },
            }
        ),
        encoding="utf-8",
    )
    settings = make_settings(
        tmp_path,
        time_series_min_candles=80,
        time_series_validation_window=20,
        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 True
    assert any(candidate["model"] == "lstm" for candidate in forecast.candidates)


def test_time_series_forecaster_reads_torch_gru_artifact(make_settings, tmp_path) -> None:
    artifact_path = tmp_path / "lstm_forecaster.json"
    hidden_size = 2
    artifact_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,
                        "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": 0.2,
                    },
                },
            }
        ),
        encoding="utf-8",
    )
    settings = make_settings(
        tmp_path,
        time_series_min_candles=80,
        time_series_validation_window=20,
        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 any(candidate["model"] == "torch_gru" for candidate in forecast.candidates)