okx-codex-trader/tests/test_run_bb_squeeze_executor.py

from dataclasses import replace

from okx_codex_trader.models import Candle
from scripts.run_bb_squeeze_executor import EMPTY_STATE, frame_from_candles, refresh_live_frame, signal_from_frame


def candles_with_latest_breakout() -> list[Candle]:
    candles = []
    price = 100.0
    for index in range(1_001):
        if index == 1_000:
            price = 101.5
        candles.append(
            Candle(
                symbol="ETH-USDT-SWAP",
                ts=1_700_000_000_000 + (index * 900_000),
                open=price,
                high=price,
                low=price,
                close=price,
                volume=1_000.0,
            )
        )
    return candles


def test_signal_uses_latest_confirmed_candle() -> None:
    candles = candles_with_latest_breakout()
    frame = frame_from_candles(candles)

    next_state, signal = signal_from_frame(frame, EMPTY_STATE)

    assert signal["decision_candle_ts"] == candles[-1].ts
    assert next_state.last_candle_ts == candles[-1].ts


def test_seen_latest_candle_replays_without_order_signal() -> None:
    candles = candles_with_latest_breakout()
    frame = frame_from_candles(candles)
    state = replace(EMPTY_STATE, last_candle_ts=candles[-1].ts)

    _, signal = signal_from_frame(frame, state)

    assert signal["decision_candle_ts"] == candles[-1].ts
    assert signal["signal"] == "state_replay"


def test_loop_predicate_skips_seen_decision_candle() -> None:
    candles = candles_with_latest_breakout()
    frame = frame_from_candles(candles)
    state = replace(EMPTY_STATE, last_candle_ts=candles[-1].ts)

    _, signal = signal_from_frame(frame, state)

    assert signal["signal"] == "state_replay"


def test_refresh_live_frame_fetches_recent_candles_after_initial_load() -> None:
    class Client:
        def __init__(self) -> None:
            self.limits = []

        def get_candles(self, symbol: str, bar: str, limit: int) -> list[Candle]:
            self.limits.append(limit)
            return candles_with_latest_breakout()

        def get_recent_candles(self, symbol: str, bar: str, limit: int) -> list[Candle]:
            self.limits.append(limit)
            candles = candles_with_latest_breakout()
            return [
                *candles[-19:],
                Candle(
                    symbol="ETH-USDT-SWAP",
                    ts=candles[-1].ts + 900_000,
                    open=102.0,
                    high=102.0,
                    low=102.0,
                    close=102.0,
                    volume=1_000.0,
                ),
            ]

    client = Client()
    initial = refresh_live_frame(client, None)
    refreshed = refresh_live_frame(client, initial)

    assert client.limits == [1_200, 20]
    assert int(refreshed.iloc[-1]["ts"]) == int(initial.iloc[-1]["ts"]) + 900_000


def middle_exit_candles(close_multiplier: float) -> list[Candle]:
    candles = []
    for index in range(1_001):
        close = 100.0
        if index == 1_000:
            close = 100.0 * close_multiplier
        candles.append(
            Candle(
                symbol="ETH-USDT-SWAP",
                ts=1_700_000_000_000 + (index * 900_000),
                open=close,
                high=close,
                low=close,
                close=close,
                volume=1_000.0,
            )
        )
    return candles


def test_short_middle_exit_requires_buffer_break() -> None:
    candles = middle_exit_candles(1.0004)
    frame = frame_from_candles(candles)
    state = replace(EMPTY_STATE, active_side="short", entry_price=101.0, entry_candle_ts=candles[-2].ts, last_candle_ts=candles[-2].ts)

    _, signal = signal_from_frame(frame, state)

    assert signal["signal"] == "hold"
    assert signal["target_side"] == "short"


def test_short_middle_exit_triggers_after_buffer_break() -> None:
    candles = middle_exit_candles(1.001)
    frame = frame_from_candles(candles)
    state = replace(EMPTY_STATE, active_side="short", entry_price=101.0, entry_candle_ts=candles[-2].ts, last_candle_ts=candles[-2].ts)

    _, signal = signal_from_frame(frame, state)

    assert signal["signal"] == "exit_middle"
    assert signal["target_side"] == "flat"