118 lines
3.4 KiB
Python
118 lines
3.4 KiB
Python
# app/tools/dzengi_probe/connection_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
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ConnectionProbeResult:
|
|
name: str
|
|
ok: bool
|
|
error: str | None = None
|
|
|
|
|
|
class DzengiConnectionProbe:
|
|
def __init__(self, config: DzengiProbeConfig) -> None:
|
|
self.config = config
|
|
|
|
async def run_all(self) -> list[ConnectionProbeResult]:
|
|
return [
|
|
await self._check(
|
|
name="baseline",
|
|
headers={
|
|
"Origin": self.config.base_url,
|
|
"Content-Type": "application/json",
|
|
},
|
|
subprotocols=[Subprotocol("json")],
|
|
),
|
|
await self._check(
|
|
name="without_origin",
|
|
headers={
|
|
"Content-Type": "application/json",
|
|
},
|
|
subprotocols=[Subprotocol("json")],
|
|
),
|
|
await self._check(
|
|
name="without_content_type",
|
|
headers={
|
|
"Origin": self.config.base_url,
|
|
},
|
|
subprotocols=[Subprotocol("json")],
|
|
),
|
|
await self._check(
|
|
name="without_headers",
|
|
headers=None,
|
|
subprotocols=[Subprotocol("json")],
|
|
),
|
|
await self._check(
|
|
name="without_subprotocol",
|
|
headers={
|
|
"Origin": self.config.base_url,
|
|
"Content-Type": "application/json",
|
|
},
|
|
subprotocols=None,
|
|
),
|
|
await self._check(
|
|
name="minimal_connection",
|
|
headers=None,
|
|
subprotocols=None,
|
|
),
|
|
]
|
|
|
|
async def _check(
|
|
self,
|
|
*,
|
|
name: str,
|
|
headers: dict[str, str] | None,
|
|
subprotocols: list[Subprotocol] | None,
|
|
) -> ConnectionProbeResult:
|
|
subscribe_message: dict[str, Any] = {
|
|
"correlationId": f"connection-probe-{name}",
|
|
"destination": "marketData.subscribe",
|
|
"payload": {
|
|
"symbols": [self.config.symbol],
|
|
},
|
|
}
|
|
|
|
try:
|
|
async with websockets.connect(
|
|
self.config.ws_url,
|
|
extra_headers=headers,
|
|
subprotocols=subprotocols,
|
|
ping_interval=20,
|
|
ping_timeout=20,
|
|
close_timeout=5,
|
|
) as websocket:
|
|
await websocket.send(json.dumps(subscribe_message))
|
|
|
|
raw_message = await asyncio.wait_for(
|
|
websocket.recv(),
|
|
timeout=self.config.timeout_seconds,
|
|
)
|
|
|
|
payload = json.loads(raw_message)
|
|
|
|
if payload.get("status") != "OK":
|
|
return ConnectionProbeResult(
|
|
name=name,
|
|
ok=False,
|
|
error=json.dumps(payload, ensure_ascii=False),
|
|
)
|
|
|
|
return ConnectionProbeResult(name=name, ok=True)
|
|
|
|
except Exception as exc:
|
|
return ConnectionProbeResult(
|
|
name=name,
|
|
ok=False,
|
|
error=str(exc),
|
|
) |