dzentra_bot/app/scripts/check_trades_unsubscribe.py

from __future__ import annotations

import asyncio
import json
import os
from typing import Any
from uuid import uuid4

import websockets
from websockets.typing import Subprotocol


SYMBOL = os.getenv(
    "TRADE_TEST_SYMBOL",
    "BTC/USD_LEVERAGE",
)

OBSERVE_BEFORE_UNSUBSCRIBE_SECONDS = 30
OBSERVE_AFTER_UNSUBSCRIBE_SECONDS = 60
RECEIVE_TIMEOUT_SECONDS = 10


def build_ws_url() -> str:
    raw_url = os.getenv(
        "EXCHANGE_WS_URL",
        "wss://api-adapter.dzengi.com",
    ).rstrip("/")

    if raw_url.startswith("https://"):
        raw_url = raw_url.replace("https://", "wss://", 1)
    elif raw_url.startswith("http://"):
        raw_url = raw_url.replace("http://", "ws://", 1)

    if raw_url.endswith("/connect"):
        return raw_url

    return f"{raw_url}/connect"


def build_headers() -> dict[str, str]:
    base_url = os.getenv(
        "EXCHANGE_BASE_URL",
        "https://api-adapter.dzengi.com",
    ).rstrip("/")

    headers = {
        "Origin": base_url,
        "Content-Type": "application/json",
    }

    api_key = os.getenv("EXCHANGE_API_KEY", "").strip()

    if api_key:
        headers["X-MBX-APIKEY"] = api_key

    return headers


def build_request(destination: str) -> dict[str, Any]:
    return {
        "correlationId": str(uuid4()),
        "destination": destination,
        "payload": {
            "symbols": [SYMBOL],
        },
    }


def decode_message(raw_message: str | bytes) -> Any:
    if isinstance(raw_message, bytes):
        raw_message = raw_message.decode("utf-8")

    try:
        return json.loads(raw_message)
    except json.JSONDecodeError:
        return raw_message


def print_message(label: str, message: Any) -> None:
    print()
    print("=" * 80)
    print(label)
    print("=" * 80)

    if isinstance(message, (dict, list)):
        print(json.dumps(message, indent=2, ensure_ascii=False))
    else:
        print(message)


def is_trade_event(message: Any) -> bool:
    return (
        isinstance(message, dict)
        and message.get("destination") == "internal.trade"
    )


async def receive_until(
    websocket: Any,
    *,
    duration_seconds: int,
    phase: str,
) -> int:
    loop = asyncio.get_running_loop()
    deadline = loop.time() + duration_seconds
    trade_count = 0

    while loop.time() < deadline:
        remaining = deadline - loop.time()
        timeout = min(RECEIVE_TIMEOUT_SECONDS, remaining)

        try:
            raw_message = await asyncio.wait_for(
                websocket.recv(),
                timeout=timeout,
            )
        except asyncio.TimeoutError:
            print(
                f"[{phase}] За последние {timeout:.0f} секунд "
                "сообщений не получено; соединение остаётся открытым."
            )
            continue

        message = decode_message(raw_message)
        print_message(f"[{phase}] Получено сообщение", message)

        if is_trade_event(message):
            trade_count += 1

    return trade_count


async def main() -> None:
    ws_url = build_ws_url()
    headers = build_headers()

    subscribe_request = build_request("trades.subscribe")
    unsubscribe_request = build_request("trades.unsubscribe")

    print(f"WebSocket URL: {ws_url}")
    print(f"Symbol: {SYMBOL}")

    async with websockets.connect(
        ws_url,
        extra_headers=headers,
        subprotocols=[Subprotocol("json")],
        ping_interval=20,
        ping_timeout=20,
        open_timeout=20,
    ) as websocket:
        print_message(
            "Отправляется trades.subscribe",
            subscribe_request,
        )

        await websocket.send(
            json.dumps(subscribe_request),
        )

        before_count = await receive_until(
            websocket,
            duration_seconds=OBSERVE_BEFORE_UNSUBSCRIBE_SECONDS,
            phase="BEFORE UNSUBSCRIBE",
        )

        print()
        print(
            "Количество Trade Event до unsubscribe: "
            f"{before_count}"
        )

        print_message(
            "Отправляется trades.unsubscribe",
            unsubscribe_request,
        )

        await websocket.send(
            json.dumps(unsubscribe_request),
        )

        after_count = await receive_until(
            websocket,
            duration_seconds=OBSERVE_AFTER_UNSUBSCRIBE_SECONDS,
            phase="AFTER UNSUBSCRIBE",
        )

        print()
        print("=" * 80)
        print("РЕЗУЛЬТАТ НАБЛЮДЕНИЯ")
        print("=" * 80)
        print(f"Trade Event до unsubscribe: {before_count}")
        print(f"Trade Event после unsubscribe: {after_count}")

        if before_count == 0:
            print(
                "До unsubscribe не было получено ни одной сделки. "
                "Эксперимент нельзя считать доказательным."
            )
        elif after_count == 0:
            print(
                "После trades.unsubscribe новые сделки не поступили. "
                "Команда, вероятно, поддерживается, но необходимо "
                "проверить ACK или повторить тест при высокой активности."
            )
        else:
            print(
                "После trades.unsubscribe сделки продолжили поступать. "
                "Команда либо не поддерживается, либо была отклонена, "
                "либо имеет другой формат."
            )


if __name__ == "__main__":
    asyncio.run(main())