okx-codex-trader/tests/test_rsi2_report.py

import pytest

from okx_codex_trader import rsi2_report
from okx_codex_trader.models import Candle
from okx_codex_trader.rsi2_report import RSI2Config, generate_rsi2_sampled_report, run_rsi2_segment
from okx_codex_trader.sampled_report import SegmentResult


def make_candle(index: int, open_price: float, close: float) -> Candle:
    high = max(open_price, close) + 1.0
    low = min(open_price, close) - 1.0
    return Candle(
        symbol="BTC-USDT-SWAP",
        ts=index * 60_000,
        open=open_price,
        high=high,
        low=low,
        close=close,
        volume=1_000.0 + index,
    )


def build_linear_candles(count: int) -> list[Candle]:
    candles: list[Candle] = []
    for index in range(count):
        price = 100.0 + index
        candles.append(make_candle(index, price, price))
    return candles


def build_long_trade_fixture() -> list[Candle]:
    return [
        make_candle(0, 90.0, 90.0),
        make_candle(1, 100.0, 100.0),
        make_candle(2, 150.0, 150.0),
        make_candle(3, 149.0, 149.0),
        make_candle(4, 148.0, 148.0),
        make_candle(5, 149.0, 149.0),
        make_candle(6, 150.0, 150.0),
    ]


def build_short_trade_fixture() -> list[Candle]:
    return [
        make_candle(0, 160.0, 160.0),
        make_candle(1, 150.0, 150.0),
        make_candle(2, 100.0, 100.0),
        make_candle(3, 101.0, 101.0),
        make_candle(4, 102.0, 102.0),
        make_candle(5, 101.0, 101.0),
        make_candle(6, 100.0, 100.0),
    ]


def build_exit_priority_fixture() -> list[Candle]:
    return [
        make_candle(0, 90.0, 90.0),
        make_candle(1, 100.0, 100.0),
        make_candle(2, 150.0, 150.0),
        make_candle(3, 149.1, 149.1),
        make_candle(4, 149.0, 149.0),
        make_candle(5, 149.1, 149.1),
        make_candle(6, 149.2, 149.2),
    ]


def build_final_bar_signal_fixture() -> list[Candle]:
    return [
        make_candle(0, 90.0, 90.0),
        make_candle(1, 100.0, 100.0),
        make_candle(2, 150.0, 150.0),
        make_candle(3, 149.0, 149.0),
        make_candle(4, 148.0, 148.0),
    ]


def build_open_tail_fixture() -> list[Candle]:
    return [
        make_candle(0, 90.0, 90.0),
        make_candle(1, 100.0, 100.0),
        make_candle(2, 150.0, 150.0),
        make_candle(3, 149.0, 149.0),
        make_candle(4, 148.0, 148.0),
        make_candle(5, 149.0, 151.0),
    ]


def build_depleted_equity_fixture() -> list[Candle]:
    return [
        make_candle(0, 90.0, 90.0),
        make_candle(1, 100.0, 100.0),
        make_candle(2, 150.0, 150.0),
        make_candle(3, 149.0, 149.0),
        make_candle(4, 148.0, 148.0),
        make_candle(5, 149.0, 151.0),
        make_candle(6, 0.0, 100.0),
        make_candle(7, 150.0, 150.0),
        make_candle(8, 149.0, 149.0),
        make_candle(9, 150.0, 150.0),
    ]


def test_compute_rsi_uses_wilder_smoothing():
    closes = rsi2_report.pd.Series([100.0, 102.07, 103.62, 103.14, 101.69], dtype=float)

    rsi = rsi2_report._compute_rsi(closes, 2)

    assert rsi[4] == pytest.approx(34.8747591522)


def test_run_rsi2_segment_produces_long_trade():
    result = run_rsi2_segment(
        candles=build_long_trade_fixture(),
        leverage=2,
        warmup_bars=4,
        config=RSI2Config(trend_sma=4, rsi_length=2, rsi_long_threshold=95.0),
    )

    assert isinstance(result, SegmentResult)
    assert result.trade_count == 1
    assert result.trades[0]["side"] == "Long"
    assert result.trades[0]["entry_price"] == pytest.approx(149.0)
    assert result.trades[0]["exit_price"] == pytest.approx(150.0)
    assert result.open_position is None


def test_run_rsi2_segment_produces_short_trade():
    result = run_rsi2_segment(
        candles=build_short_trade_fixture(),
        leverage=2,
        warmup_bars=4,
        config=RSI2Config(trend_sma=4, rsi_length=2, rsi_short_threshold=5.0),
    )

    assert isinstance(result, SegmentResult)
    assert result.trade_count == 1
    assert result.trades[0]["side"] == "Short"
    assert result.trades[0]["entry_price"] == pytest.approx(101.0)
    assert result.trades[0]["exit_price"] == pytest.approx(100.0)
    assert result.open_position is None


def test_run_rsi2_segment_exit_priority_is_correct():
    result = run_rsi2_segment(
        candles=build_exit_priority_fixture(),
        leverage=2,
        warmup_bars=4,
        config=RSI2Config(trend_sma=4, rsi_length=2, rsi_long_threshold=98.0, rsi_short_threshold=50.0, exit_rsi=96.5),
    )

    assert result.trade_count == 1
    assert len(result.entries) == 1
    assert result.trades[0]["side"] == "Long"
    assert result.open_position is None


def test_run_rsi2_segment_does_not_generate_entry_from_final_bar():
    result = run_rsi2_segment(
        candles=build_final_bar_signal_fixture(),
        leverage=2,
        warmup_bars=4,
        config=RSI2Config(trend_sma=4, rsi_length=2, rsi_long_threshold=95.0),
    )

    assert result.trade_count == 0
    assert result.entries == []
    assert result.open_position is None


def test_run_rsi2_segment_marks_open_position_to_market():
    result = run_rsi2_segment(
        candles=build_open_tail_fixture(),
        leverage=2,
        warmup_bars=4,
        config=RSI2Config(trend_sma=4, rsi_length=2, rsi_long_threshold=95.0),
    )

    assert result.trade_count == 0
    assert result.trades == []
    assert result.total_return == pytest.approx((10_268.456375838927 - 10_000.0) / 10_000.0)
    assert result.open_position is not None
    assert result.open_position["side"] == "long"


def test_run_rsi2_segment_stops_after_equity_is_depleted():
    result = run_rsi2_segment(
        candles=build_depleted_equity_fixture(),
        leverage=2,
        warmup_bars=4,
        config=RSI2Config(trend_sma=4, rsi_length=2, rsi_long_threshold=95.0),
    )

    assert result.trade_count == 1
    assert result.open_position is None
    assert len(result.entries) == 1
    assert result.total_return <= -1.0


def test_generate_rsi2_sampled_report_uses_shared_shell_defaults(monkeypatch, tmp_path):
    candles = build_linear_candles(5_000)
    output_file = tmp_path / "rsi2.html"
    recorded: dict[str, object] = {}
    sentinel = {
        "report_file": str(output_file),
        "segment_count": 2,
        "window_size": 300,
        "aggregate_trade_count": 4,
        "average_return": 0.12,
    }

    def fake_generate_sampled_report(**kwargs):
        recorded.update(kwargs)
        return sentinel

    monkeypatch.setattr(rsi2_report, "generate_sampled_report", fake_generate_sampled_report)

    result = generate_rsi2_sampled_report(
        candles=candles,
        leverage=2,
        output_file=output_file,
        symbol="BTC-USDT-SWAP",
        bar="3m",
        segments=2,
        window_size=300,
    )

    assert result == sentinel
    assert recorded["report_title"] == "RSI2 Sampled Report"
    assert recorded["strategy_label"] == "RSI2"
    assert recorded["strategy_params"] == {
        "trend_sma": 50,
        "rsi_length": 2,
        "rsi_long_threshold": 10.0,
        "rsi_short_threshold": 90.0,
        "exit_rsi": 50.0,
    }
    assert recorded["warmup_bars"] == 50
    assert callable(recorded["run_segment"])