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

1""" 

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

3 

4Module 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 

26import threading 

27import json 

28import websocket 

29import logging 

30from threading import Condition 

31from collections.abc import Callable 

32from typing import Any 

33 

34 

35logger = logging.getLogger(__name__) 

36 

37 

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