Coverage for src/python/dandeliion/client/websocket.py: 97%

1"""

2@file src/python/dandeliion/client/websocket.py

3

4 Module for websocket client used in simulator

5"""

6

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#

24

25# built-in modules

26 import threading

27 import json

28 import websocket

29 import logging

30 from threading import Condition

31 from collections.abc import Callable

32 from typing import Any

33

34

35 logger = logging.getLogger(__name__)

36

37

38 class SimulatorWebSocketClient:

39

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

42

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 )

51

52 self._on_update = on_update

53 self._is_opened = False

54 self._is_ready = Condition()

55

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()

60

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)

65

66 def subscribe(self, run_id):

67 self.send_message(run_id)

68

69 def close(self):

70 self._app.close()

71

72 def _on_open(self, wsapp):

73 with self._is_ready:

74 self._is_opened = True

75 self._is_ready.notify_all()

76

77 def _on_message(self, wsapp, message) -> None:

78 self._on_update(json.loads(message)['updates'])

79

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))

85

86 def _on_error(self, wsapp, err):

87 logger.error("ERROR", wsapp, err)