Coverage for src/python/dandeliion/client/websocket.py: 97%
37 statements
« prev ^ index » next coverage.py v7.6.12, created at 2025-02-23 23:59 +0100
« prev ^ index » next coverage.py v7.6.12, created at 2025-02-23 23:59 +0100
1"""
2@file src/python/dandeliion/client/websocket.py
4Module for websocket client used in simulator
5"""
7#
8# Copyright (C) 2024-2025 Dandeliion Team
9#
10# This library is free software; you can redistribute it and/or modify it under
11# the terms of the GNU Lesser General Public License as published by the Free
12# Software Foundation; either version 3.0 of the License, or (at your option)
13# any later version.
14#
15# This library is distributed in the hope that it will be useful, but WITHOUT
16# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
17# FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public License for more
18# details.
19#
20# You should have received a copy of the GNU Lesser General Public License
21# along with this library; if not, write to the Free Software Foundation, Inc.,
22# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
23#
25# built-in modules
26import threading
27import json
28import websocket
29import logging
30from threading import Condition
31from collections.abc import Callable
32from typing import Any
35logger = logging.getLogger(__name__)
38class SimulatorWebSocketClient:
40 def __init__(self, url: str, api_key: str, on_update: Callable[[Any], None]):
41 headers = {'Authorization': f'Token {api_key}'} # TODO adapt to server
43 self._app = websocket.WebSocketApp(
44 url,
45 on_open=self._on_open,
46 on_message=self._on_message,
47 on_error=self._on_error,
48 on_close=self._on_close,
49 header=headers,
50 )
52 self._on_update = on_update
53 self._is_opened = False
54 self._is_ready = Condition()
56 # Initialise the run_forever inside a thread and make this thread as a daemon thread
57 wst = threading.Thread(target=self._app.run_forever)
58 # wst.daemon = True # not needed anymore?
59 wst.start()
61 def send_message(self, message):
62 with self._is_ready:
63 self._is_ready.wait_for(lambda: self._is_opened)
64 self._app.send(message)
66 def subscribe(self, run_id):
67 self.send_message(run_id)
69 def close(self):
70 self._app.close()
72 def _on_open(self, wsapp):
73 with self._is_ready:
74 self._is_opened = True
75 self._is_ready.notify_all()
77 def _on_message(self, wsapp, message) -> None:
78 self._on_update(json.loads(message)['updates'])
80 def _on_close(self, wsapp, close_status_code, close_msg):
81 # Because on_close was triggered, we know the opcode = 8
82 logger.debug("on_close args:")
83 logger.debug("close status code: " + str(close_status_code))
84 logger.debug("close message: " + str(close_msg))
86 def _on_error(self, wsapp, err):
87 logger.error("ERROR", wsapp, err)