134 lines
4.0 KiB
Python
134 lines
4.0 KiB
Python
# app/tools/dzengi_probe/websocket_probe.py
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
import websockets
|
|
from websockets.typing import Subprotocol
|
|
|
|
from app.tools.dzengi_probe.config import DzengiProbeConfig
|
|
from app.tools.dzengi_probe.response_store import append_jsonl, save_json
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class WebSocketProbeResult:
|
|
stream: str
|
|
output_file: str
|
|
ok: bool
|
|
messages_saved: int
|
|
error: str | None = None
|
|
|
|
|
|
class DzengiWebSocketProbe:
|
|
def __init__(self, config: DzengiProbeConfig) -> None:
|
|
self.config = config
|
|
|
|
async def probe_market_data(self, *, max_messages: int = 10) -> WebSocketProbeResult:
|
|
output_file = (
|
|
f"websocket/marketData.subscribe/"
|
|
f"{self.config.symbol_file_name}.jsonl"
|
|
)
|
|
|
|
request_file = (
|
|
f"websocket/marketData.subscribe/"
|
|
f"{self.config.symbol_file_name}.request.json"
|
|
)
|
|
|
|
subscribe_message: dict[str, Any] = {
|
|
"correlationId": "probe-market-data-subscribe",
|
|
"destination": "marketData.subscribe",
|
|
"payload": {
|
|
"symbols": [self.config.symbol],
|
|
},
|
|
}
|
|
|
|
save_json(self.config.output_dir / request_file, subscribe_message)
|
|
|
|
return await self._subscribe_and_save(
|
|
stream="marketData.subscribe",
|
|
subscribe_message=subscribe_message,
|
|
output_file=output_file,
|
|
max_messages=max_messages,
|
|
)
|
|
|
|
async def _subscribe_and_save(
|
|
self,
|
|
*,
|
|
stream: str,
|
|
subscribe_message: dict[str, Any],
|
|
output_file: str,
|
|
max_messages: int,
|
|
) -> WebSocketProbeResult:
|
|
output_path = self.config.output_dir / output_file
|
|
messages_saved = 0
|
|
|
|
if output_path.exists():
|
|
output_path.unlink()
|
|
|
|
try:
|
|
async with websockets.connect(
|
|
self.config.ws_url,
|
|
extra_headers={
|
|
"Origin": self.config.base_url,
|
|
"Content-Type": "application/json",
|
|
},
|
|
subprotocols=[Subprotocol("json")],
|
|
ping_interval=20,
|
|
ping_timeout=20,
|
|
close_timeout=5,
|
|
) as websocket:
|
|
await websocket.send(json.dumps(subscribe_message))
|
|
|
|
while messages_saved < max_messages:
|
|
raw_message = await asyncio.wait_for(
|
|
websocket.recv(),
|
|
timeout=self.config.timeout_seconds,
|
|
)
|
|
|
|
try:
|
|
payload = json.loads(raw_message)
|
|
except json.JSONDecodeError:
|
|
payload = {"raw": raw_message}
|
|
|
|
print(payload)
|
|
append_jsonl(output_path, payload)
|
|
messages_saved += 1
|
|
|
|
return WebSocketProbeResult(
|
|
stream=stream,
|
|
output_file=str(output_path),
|
|
ok=True,
|
|
messages_saved=messages_saved,
|
|
)
|
|
|
|
except Exception as exc:
|
|
return WebSocketProbeResult(
|
|
stream=stream,
|
|
output_file=str(output_path),
|
|
ok=False,
|
|
messages_saved=messages_saved,
|
|
error=str(exc),
|
|
)
|
|
|
|
async def probe_depth_request(self) -> WebSocketProbeResult:
|
|
output_file = f"websocket/depth/{self.config.symbol_file_name}.jsonl"
|
|
|
|
request_message: dict[str, Any] = {
|
|
"correlationId": "probe-depth-request",
|
|
"destination": f"/api/{self.config.api_version}/depth",
|
|
"payload": {
|
|
"symbol": self.config.symbol,
|
|
"limit": int(self.config.depth_limit),
|
|
},
|
|
}
|
|
|
|
return await self._subscribe_and_save(
|
|
stream="wss:/api/v1/depth",
|
|
subscribe_message=request_message,
|
|
output_file=output_file,
|
|
max_messages=1,
|
|
) |