72 lines
2.4 KiB
Python
72 lines
2.4 KiB
Python
import json
|
|
import time
|
|
from src.libs.common import write_signal_to_file
|
|
|
|
import websocket
|
|
|
|
from src.libs.custom_logger import get_custom_logger
|
|
|
|
logger = get_custom_logger(__name__)
|
|
|
|
|
|
class SignalClient:
|
|
|
|
def __init__(self, ws_url, await_signals):
|
|
self.url = f"{ws_url}/signals"
|
|
self.await_signals = await_signals
|
|
self.received_signals = {signal: [] for signal in self.await_signals}
|
|
|
|
def on_message(self, ws, signal):
|
|
signal_data = json.loads(signal)
|
|
signal_type = signal_data.get("type")
|
|
|
|
write_signal_to_file(signal_data)
|
|
|
|
if signal_type in self.await_signals:
|
|
self.received_signals[signal_type].append(signal_data)
|
|
|
|
def wait_for_signal(self, signal_type, timeout=20):
|
|
start_time = time.time()
|
|
while not self.received_signals.get(signal_type):
|
|
if time.time() - start_time >= timeout:
|
|
raise TimeoutError(
|
|
f"Signal {signal_type} is not received in {timeout} seconds")
|
|
time.sleep(0.2)
|
|
logger.debug(f"Signal {signal_type} is received in {round(time.time() - start_time)} seconds")
|
|
return self.received_signals[signal_type][0]
|
|
|
|
def wait_for_complete_signal(self, signal_type, timeout=5):
|
|
start_time = time.time()
|
|
events = []
|
|
|
|
while time.time() - start_time < timeout:
|
|
if self.received_signals.get(signal_type):
|
|
events.extend(self.received_signals[signal_type])
|
|
self.received_signals[signal_type] = []
|
|
time.sleep(0.2)
|
|
|
|
if events:
|
|
logger.debug(
|
|
f"Collected {len(events)} events of type {signal_type} within {timeout} seconds")
|
|
return events
|
|
raise TimeoutError(f"No signals of type {signal_type} received in {timeout} seconds")
|
|
|
|
def _on_error(self, ws, error):
|
|
logger.error(f"WebSocket error: {error}")
|
|
|
|
def _on_close(self, ws, close_status_code, close_msg):
|
|
logger.info(f"WebSocket connection closed: {close_status_code}, {close_msg}")
|
|
|
|
def _on_open(self, ws):
|
|
logger.info("WebSocket connection opened")
|
|
|
|
def _connect(self):
|
|
ws = websocket.WebSocketApp(
|
|
self.url,
|
|
on_message=self.on_message,
|
|
on_error=self._on_error,
|
|
on_close=self._on_close
|
|
)
|
|
ws.on_open = self._on_open
|
|
ws.run_forever()
|