okx-codex-trader/tests/test_okx_client.py

import base64
import hashlib
import hmac
from dataclasses import dataclass
from urllib.parse import urlencode, urlparse

import pytest

from okx_codex_trader.config import Config
from okx_codex_trader.models import InstrumentMeta, TradeSignal
from okx_codex_trader.okx_client import OkxClient, build_contract_size


@dataclass
class DummyResponse:
    payload: dict[str, object]
    status_code: int = 200
    json_error: Exception | None = None

    def json(self) -> dict[str, object]:
        if self.json_error is not None:
            raise self.json_error
        return self.payload


@dataclass
class RecordedRequest:
    method: str
    url: str
    headers: dict[str, str]
    params: dict[str, object] | None
    json_body: dict[str, object] | None


class DummySession:
    def __init__(self, responses: list[DummyResponse | Exception] | None = None):
        self._responses = list(responses or [])
        self.last_request: RecordedRequest | None = None
        self.request_paths: list[str] = []
        self.request_bodies: list[dict[str, object] | None] = []

    @property
    def last_json_body(self) -> dict[str, object] | None:
        return self.last_request.json_body if self.last_request else None

    def request(
        self,
        method: str,
        url: str,
        *,
        headers: dict[str, str] | None = None,
        params: dict[str, object] | None = None,
        json: dict[str, object] | None = None,
    ) -> DummyResponse:
        self.last_request = RecordedRequest(
            method=method,
            url=url,
            headers=headers or {},
            params=params,
            json_body=json,
        )
        self.request_paths.append(urlparse(url).path)
        self.request_bodies.append(json)
        if self._responses:
            response = self._responses.pop(0)
            if isinstance(response, Exception):
                raise response
            return response
        return candles_response()


def sample_config() -> Config:
    return Config(api_key="key", api_secret="secret", api_passphrase="passphrase")


def candles_response() -> DummyResponse:
    return DummyResponse(
        {
            "code": "0",
            "msg": "",
            "data": [
                ["1710000000000", "25000", "25100", "24900", "25050", "100", "1000", "1000", "1"],
            ],
        }
    )


def descending_candles_response() -> DummyResponse:
    return DummyResponse(
        {
            "code": "0",
            "msg": "",
            "data": [
                ["1710000001000", "25100", "25200", "25000", "25150", "110", "1100", "1100", "1"],
                ["1710000000000", "25000", "25100", "24900", "25050", "100", "1000", "1000", "1"],
            ],
        }
    )


def instrument_response() -> DummyResponse:
    return DummyResponse(
        {
            "code": "0",
            "msg": "",
            "data": [
                {
                    "instId": "BTC-USDT-SWAP",
                    "instType": "SWAP",
                    "ctVal": "0.001",
                    "lotSz": "1",
                    "minSz": "1",
                }
            ],
        }
    )


def large_min_size_instrument_response() -> DummyResponse:
    return DummyResponse(
        {
            "code": "0",
            "msg": "",
            "data": [
                {
                    "instId": "BTC-USDT-SWAP",
                    "instType": "SWAP",
                    "ctVal": "0.01",
                    "lotSz": "1",
                    "minSz": "100",
                }
            ],
        }
    )


def ticker_response(last: str) -> DummyResponse:
    return DummyResponse({"code": "0", "msg": "", "data": [{"instId": "BTC-USDT-SWAP", "last": last}]})


def account_config_response(pos_mode: str) -> DummyResponse:
    return DummyResponse({"code": "0", "msg": "", "data": [{"posMode": pos_mode}]})


def leverage_response() -> DummyResponse:
    return DummyResponse({"code": "0", "msg": "", "data": [{"lever": "2"}]})


def place_order_response() -> DummyResponse:
    return DummyResponse({"code": "0", "msg": "", "data": [{"ordId": "123"}]})


def place_order_response_without_order_id() -> DummyResponse:
    return DummyResponse({"code": "0", "msg": "", "data": [{}]})


def error_response(code: str, msg: str) -> DummyResponse:
    return DummyResponse({"code": code, "msg": msg, "data": []})


def positions_response() -> DummyResponse:
    return DummyResponse(
        {
            "code": "0",
            "msg": "",
            "data": [
                {
                    "instId": "BTC-USDT-SWAP",
                    "posSide": "long",
                    "pos": "8",
                    "avgPx": "25000",
                }
            ],
        }
    )


def positions_with_zero_size_response() -> DummyResponse:
    return DummyResponse(
        {
            "code": "0",
            "msg": "",
            "data": [
                {
                    "instId": "BTC-USDT-SWAP",
                    "posSide": "long",
                    "pos": "0",
                    "avgPx": "25000",
                },
                {
                    "instId": "BTC-USDT-SWAP",
                    "posSide": "short",
                    "pos": "3",
                    "avgPx": "24900",
                },
            ],
        }
    )


def positions_with_zero_size_malformed_avg_price_response() -> DummyResponse:
    return DummyResponse(
        {
            "code": "0",
            "msg": "",
            "data": [
                {
                    "instId": "BTC-USDT-SWAP",
                    "posSide": "long",
                    "pos": "0",
                    "avgPx": "bad",
                },
                {
                    "instId": "BTC-USDT-SWAP",
                    "posSide": "short",
                    "pos": "3",
                    "avgPx": "24900",
                },
            ],
        }
    )


def market_long_signal() -> TradeSignal:
    return TradeSignal(
        action="long",
        confidence=0.9,
        leverage=2,
        entry_price=None,
        take_profit_price=26000.0,
        stop_loss_price=24000.0,
        reason="trend",
    )


def limit_short_signal() -> TradeSignal:
    return TradeSignal(
        action="short",
        confidence=0.8,
        leverage=2,
        entry_price=25000.0,
        take_profit_price=24000.0,
        stop_loss_price=25500.0,
        reason="mean reversion",
    )


def flat_signal() -> TradeSignal:
    return TradeSignal(
        action="flat",
        confidence=0.7,
        leverage=2,
        entry_price=None,
        take_profit_price=None,
        stop_loss_price=None,
        reason="exit",
    )


def test_signed_demo_request_attaches_headers():
    session = DummySession()
    client = OkxClient(config=sample_config(), session=session)

    client.get_candles(symbol="BTC-USDT-SWAP", bar="1H", limit=20)

    request = session.last_request
    assert request is not None
    assert request.headers["x-simulated-trading"] == "1"
    assert request.headers["OK-ACCESS-KEY"] == "key"
    assert request.headers["OK-ACCESS-PASSPHRASE"] == "passphrase"
    timestamp = request.headers["OK-ACCESS-TIMESTAMP"]
    path = urlparse(request.url).path
    query = urlencode(request.params or {})
    path_with_query = path if not query else f"{path}?{query}"
    expected_signature = base64.b64encode(
        hmac.new(
            b"secret",
            f"{timestamp}{request.method}{path_with_query}".encode(),
            hashlib.sha256,
        ).digest()
    ).decode()
    assert request.headers["OK-ACCESS-SIGN"] == expected_signature


def test_get_candles_returns_chronological_ascending_order():
    session = DummySession([descending_candles_response()])
    client = OkxClient(config=sample_config(), session=session)

    candles = client.get_candles(symbol="BTC-USDT-SWAP", bar="1H", limit=20)

    assert [candle.ts for candle in candles] == [1710000000000, 1710000001000]


def test_build_contract_size_rounds_down_to_lot_size():
    metadata = InstrumentMeta(ct_val=0.01, lot_sz=0.1, min_sz=0.1)
    assert build_contract_size(notional=251, price=25_000, metadata=metadata) == 1.0


def test_build_contract_size_fails_below_min_size():
    metadata = InstrumentMeta(ct_val=0.01, lot_sz=1, min_sz=5)
    with pytest.raises(ValueError):
        build_contract_size(notional=250, price=25_100, metadata=metadata)


def test_market_order_fetches_latest_price_before_sizing():
    session = DummySession(
        [
            instrument_response(),
            ticker_response(last="25000"),
            account_config_response(pos_mode="long_short_mode"),
            leverage_response(),
            place_order_response(),
        ]
    )
    client = OkxClient(config=sample_config(), session=session)

    client.place_demo_order(symbol="BTC-USDT-SWAP", signal=market_long_signal(), margin_usdt=100)

    assert session.request_paths == [
        "/api/v5/public/instruments",
        "/api/v5/market/ticker",
        "/api/v5/account/config",
        "/api/v5/account/set-leverage",
        "/api/v5/trade/order",
    ]


def test_place_demo_order_fails_when_not_hedge_mode():
    session = DummySession(
        [
            instrument_response(),
            ticker_response(last="25000"),
            account_config_response(pos_mode="net_mode"),
        ]
    )
    client = OkxClient(config=sample_config(), session=session)

    with pytest.raises(ValueError):
        client.place_demo_order(symbol="BTC-USDT-SWAP", signal=market_long_signal(), margin_usdt=100)


def test_place_demo_order_validates_size_before_setting_leverage():
    session = DummySession(
        [
            large_min_size_instrument_response(),
            ticker_response(last="25000"),
            account_config_response(pos_mode="long_short_mode"),
        ]
    )
    client = OkxClient(config=sample_config(), session=session)

    with pytest.raises(ValueError, match="contract size below minimum"):
        client.place_demo_order(symbol="BTC-USDT-SWAP", signal=market_long_signal(), margin_usdt=100)

    assert session.request_paths == [
        "/api/v5/public/instruments",
        "/api/v5/market/ticker",
        "/api/v5/account/config",
    ]


def test_limit_short_order_uses_sell_and_short_pos_side():
    session = DummySession(
        [
            instrument_response(),
            account_config_response(pos_mode="long_short_mode"),
            leverage_response(),
            place_order_response(),
        ]
    )
    client = OkxClient(config=sample_config(), session=session)

    client.place_demo_order(symbol="ETH-USDT-SWAP", signal=limit_short_signal(), margin_usdt=100)

    order_request = session.last_json_body
    assert order_request is not None
    assert order_request["ordType"] == "limit"
    assert order_request["side"] == "sell"
    assert order_request["posSide"] == "short"
    assert order_request["px"] == "25000"
    assert session.request_bodies[2]["lever"] == "2"
    assert session.request_bodies[2]["mgnMode"] == "isolated"


def test_flat_signal_returns_noop_without_order_submission():
    session = DummySession([])
    client = OkxClient(config=sample_config(), session=session)

    result = client.place_demo_order(symbol="BTC-USDT-SWAP", signal=flat_signal(), margin_usdt=100)

    assert result.status == "noop"
    assert session.request_paths == []


def test_place_demo_order_sends_computed_sz_and_ignores_tp_sl_fields():
    session = DummySession(
        [
            instrument_response(),
            ticker_response(last="25000"),
            account_config_response(pos_mode="long_short_mode"),
            leverage_response(),
            place_order_response(),
        ]
    )
    client = OkxClient(config=sample_config(), session=session)

    client.place_demo_order(symbol="BTC-USDT-SWAP", signal=market_long_signal(), margin_usdt=100)

    order_request = session.last_json_body
    assert order_request is not None
    assert order_request["sz"] == "8"
    assert "tpTriggerPx" not in order_request
    assert "slTriggerPx" not in order_request


def test_okx_error_payload_raises_value_error():
    session = DummySession([error_response(code="51000", msg="parameter error")])
    client = OkxClient(config=sample_config(), session=session)

    with pytest.raises(ValueError):
        client.get_candles(symbol="BTC-USDT-SWAP", bar="1H", limit=20)


def test_transport_failure_raises_stable_value_error():
    session = DummySession([RuntimeError("socket closed")])
    client = OkxClient(config=sample_config(), session=session)

    with pytest.raises(ValueError, match="okx transport error"):
        client.get_candles(symbol="BTC-USDT-SWAP", bar="1H", limit=20)


def test_invalid_json_raises_stable_value_error():
    session = DummySession([DummyResponse({}, json_error=ValueError("bad json"))])
    client = OkxClient(config=sample_config(), session=session)

    with pytest.raises(ValueError, match="okx response payload is invalid"):
        client.get_candles(symbol="BTC-USDT-SWAP", bar="1H", limit=20)


def test_empty_positions_data_returns_empty_list():
    session = DummySession([DummyResponse({"code": "0", "msg": "", "data": []})])
    client = OkxClient(config=sample_config(), session=session)

    assert client.get_positions(symbol="BTC-USDT-SWAP") == []


def test_malformed_numeric_field_raises_stable_value_error():
    session = DummySession(
        [
            DummyResponse(
                {
                    "code": "0",
                    "msg": "",
                    "data": [
                        {
                            "instId": "BTC-USDT-SWAP",
                            "posSide": "long",
                            "pos": "bad",
                            "avgPx": "25000",
                        }
                    ],
                }
            )
        ]
    )
    client = OkxClient(config=sample_config(), session=session)

    with pytest.raises(ValueError, match="okx response payload is invalid"):
        client.get_positions(symbol="BTC-USDT-SWAP")


def test_non_list_okx_data_raises_stable_value_error():
    session = DummySession([DummyResponse({"code": "0", "msg": "", "data": {}})])
    client = OkxClient(config=sample_config(), session=session)

    with pytest.raises(ValueError, match="okx response payload is invalid"):
        client.get_positions(symbol="BTC-USDT-SWAP")


def test_place_demo_order_raises_when_order_id_is_missing():
    session = DummySession(
        [
            instrument_response(),
            ticker_response(last="25000"),
            account_config_response(pos_mode="long_short_mode"),
            leverage_response(),
            place_order_response_without_order_id(),
        ]
    )
    client = OkxClient(config=sample_config(), session=session)

    with pytest.raises(ValueError, match="okx response payload is invalid"):
        client.place_demo_order(symbol="BTC-USDT-SWAP", signal=market_long_signal(), margin_usdt=100)


def test_place_demo_order_rejects_invalid_leverage_before_okx():
    session = DummySession([])
    signal = TradeSignal(
        action="long",
        confidence=0.9,
        leverage=4,
        entry_price=None,
        take_profit_price=None,
        stop_loss_price=None,
        reason="x",
    )
    client = OkxClient(config=sample_config(), session=session)

    with pytest.raises(ValueError, match="leverage is invalid"):
        client.place_demo_order(symbol="BTC-USDT-SWAP", signal=signal, margin_usdt=100)
    assert session.request_paths == []


def test_place_demo_order_rejects_unknown_action_before_okx():
    session = DummySession([])
    signal = TradeSignal(
        action="hold",
        confidence=0.9,
        leverage=2,
        entry_price=None,
        take_profit_price=None,
        stop_loss_price=None,
        reason="x",
    )
    client = OkxClient(config=sample_config(), session=session)

    with pytest.raises(ValueError, match="action is invalid"):
        client.place_demo_order(symbol="BTC-USDT-SWAP", signal=signal, margin_usdt=100)
    assert session.request_paths == []


def test_get_positions_returns_normalized_positions():
    session = DummySession([positions_response()])
    client = OkxClient(config=sample_config(), session=session)

    positions = client.get_positions(symbol="BTC-USDT-SWAP")

    assert positions[0].symbol == "BTC-USDT-SWAP"
    assert positions[0].pos_side == "long"
    assert positions[0].size == 8.0
    assert positions[0].avg_price == 25000.0


def test_get_positions_filters_zero_size_rows():
    session = DummySession([positions_with_zero_size_response()])
    client = OkxClient(config=sample_config(), session=session)

    positions = client.get_positions(symbol="BTC-USDT-SWAP")

    assert len(positions) == 1
    assert positions[0].pos_side == "short"
    assert positions[0].size == 3.0


def test_get_positions_ignores_malformed_fields_on_zero_size_rows():
    session = DummySession([positions_with_zero_size_malformed_avg_price_response()])
    client = OkxClient(config=sample_config(), session=session)

    positions = client.get_positions(symbol="BTC-USDT-SWAP")

    assert len(positions) == 1
    assert positions[0].pos_side == "short"
    assert positions[0].avg_price == 24900.0