llm-trade-bot/tests/test_bybit_gateway.py

from __future__ import annotations

from decimal import Decimal
from typing import Any

from llm_trade_bot.bybit_gateway import BybitGateway
from llm_trade_bot.config import BybitConfig, ScannerConfig
from llm_trade_bot.models import Category, OpenTradeRequest, OrderType, Side, ValidatedOrder


class FakeBybitHTTP:
    def __init__(
        self,
        wallet: dict[str, Any] | None = None,
        open_orders: dict[str, list[dict[str, Any]]] | None = None,
    ) -> None:
        self.wallet = wallet
        self.open_orders = open_orders or {}
        self.open_order_requests: list[dict[str, Any]] = []
        self.placed_orders: list[dict[str, Any]] = []

    def get_wallet_balance(self, **parameters: Any) -> dict[str, Any]:
        assert parameters == {"accountType": "UNIFIED"}
        return self.wallet or {"retCode": 0, "result": {"list": []}}

    def get_positions(self, **parameters: Any) -> dict[str, Any]:
        return {"retCode": 0, "result": {"list": []}}

    def get_open_orders(self, **parameters: Any) -> dict[str, Any]:
        self.open_order_requests.append(parameters)
        return {
            "retCode": 0,
            "result": {"list": self.open_orders.get(str(parameters["category"]), [])},
        }

    def place_order(self, **parameters: Any) -> dict[str, Any]:
        self.placed_orders.append(parameters)
        return {"retCode": 0, "result": {"orderId": "order-1"}}


async def test_open_orders_scopes_linear_requests_by_settle_coin() -> None:
    gateway = BybitGateway(BybitConfig("key", "secret", "USDC"), ScannerConfig())
    http = FakeBybitHTTP()
    gateway._http = http  # type: ignore[assignment]

    await gateway.open_orders()

    assert http.open_order_requests == [
        {"category": "spot", "limit": 50},
        {"category": "linear", "limit": 50, "settleCoin": "USDC"},
    ]


async def test_account_snapshot_uses_quote_coin_balance_when_account_available_is_empty() -> None:
    wallet = {
        "retCode": 0,
        "result": {
            "list": [
                {
                    "totalEquity": "100000",
                    "totalAvailableBalance": "99999",
                    "totalMarginBalance": "100000",
                    "coin": [
                        {
                            "coin": "USDT",
                            "walletBalance": "50000",
                            "availableToWithdraw": "",
                            "totalPositionIM": "100",
                            "totalOrderIM": "200",
                            "locked": "300",
                            "bonus": "400",
                        },
                        {"coin": "USDC", "walletBalance": "50000", "availableToWithdraw": ""},
                    ],
                }
            ]
        },
    }
    gateway = BybitGateway(BybitConfig("key", "secret", "USDC"), ScannerConfig(quote_coin="USDT"))
    gateway._http = FakeBybitHTTP(wallet)  # type: ignore[assignment]

    spot = await gateway.account_snapshot("spot")
    linear = await gateway.account_snapshot("linear")

    assert spot.available_balance == Decimal("49000")
    assert linear.available_balance == Decimal("50000")


async def test_account_snapshot_uses_open_spot_orders_when_reported_locked_is_stale() -> None:
    wallet = {
        "retCode": 0,
        "result": {
            "list": [
                {
                    "totalEquity": "100005",
                    "totalAvailableBalance": "50005",
                    "coin": [
                        {
                            "coin": "USDT",
                            "walletBalance": "100005",
                            "availableToWithdraw": "",
                            "totalPositionIM": "0",
                            "totalOrderIM": "50000",
                            "locked": "0",
                            "bonus": "0",
                        }
                    ],
                }
            ]
        },
    }
    open_orders = {
        "spot": [
            {
                "category": "spot",
                "symbol": "BTCUSDT",
                "side": "Buy",
                "orderStatus": "New",
                "leavesQty": "0.586166",
                "leavesValue": "49999.96",
                "price": "85300",
            }
        ]
    }
    gateway = BybitGateway(BybitConfig("key", "secret", "USDT"), ScannerConfig(quote_coin="USDT"))
    gateway._http = FakeBybitHTTP(wallet, open_orders)  # type: ignore[assignment]

    snapshot = await gateway.account_snapshot("linear")

    assert snapshot.available_balance == Decimal("5.04")


async def test_spot_market_order_sends_base_coin_quantity() -> None:
    gateway = BybitGateway(BybitConfig("key", "secret"), ScannerConfig())
    http = FakeBybitHTTP()
    gateway._http = http  # type: ignore[assignment]
    request = OpenTradeRequest(
        category=Category.SPOT,
        symbol="BTCUSDT",
        side=Side.BUY,
        order_type=OrderType.MARKET,
        reference_price=Decimal("100000"),
        limit_price=None,
        stop_loss=Decimal("90000"),
        take_profit=Decimal("110000"),
        trailing_distance=None,
        leverage=Decimal(1),
        rationale="test",
    )
    order = ValidatedOrder(
        entry_price=Decimal("100000"),
        quantity=Decimal("0.5"),
        notional=Decimal("50000"),
        stop_loss=Decimal("90000"),
        take_profit=Decimal("110000"),
        trailing_distance=None,
        leverage=Decimal(1),
    )

    await gateway.open_trade(request, order, "ltb-test")

    assert http.placed_orders[0]["qty"] == "0.5"
    assert http.placed_orders[0]["marketUnit"] == "baseCoin"