| # |
| # Copyright (c) 2023 Project CHIP Authors |
| # |
| # Licensed under the Apache License, Version 2.0 (the 'License'); |
| # you may not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an 'AS IS' BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| |
| import logging |
| import re |
| import select |
| import subprocess |
| import time |
| from dataclasses import dataclass |
| |
| import websockets |
| |
| from .hooks import WebSocketRunnerHooks |
| from .runner import TestRunner |
| |
| _KEEP_ALIVE_TIMEOUT_IN_SECONDS = 120 |
| _MAX_MESSAGE_SIZE_IN_BYTES = 10485760 # 10 MB |
| _CONNECT_MAX_RETRIES_DEFAULT = 4 |
| _WEBSOCKET_SERVER_MESSAGE = '== WebSocket Server Ready' |
| _WEBSOCKET_SERVER_MESSAGE_TIMEOUT = 60 # seconds |
| _WEBSOCKET_SERVER_TERMINATE_TIMEOUT = 10 # seconds |
| |
| |
| @dataclass |
| class WebSocketRunnerConfig: |
| server_address: str = 'localhost' |
| server_port: int = '9002' |
| server_path: str = None |
| server_arguments: str = None |
| hooks: WebSocketRunnerHooks = WebSocketRunnerHooks() |
| |
| |
| class WebSocketRunner(TestRunner): |
| def __init__(self, config: WebSocketRunnerConfig): |
| self._client = None |
| self._server = None |
| self._hooks = config.hooks |
| |
| self._server_connection_url = self._make_server_connection_url( |
| config.server_address, config.server_port) |
| self._server_startup_command = self._make_server_startup_command( |
| config.server_path, config.server_arguments, config.server_port) |
| |
| @property |
| def is_connected(self) -> bool: |
| if self._client is None: |
| return False |
| |
| return self._client.state == websockets.protocol.State.OPEN |
| |
| async def start(self): |
| self._server = await self._start_server(self._server_startup_command, self._server_connection_url) |
| self._client = await self._start_client(self._server_connection_url) |
| |
| async def stop(self): |
| await self._stop_client(self._client) |
| await self._stop_server(self._server) |
| self._client = None |
| self._server = None |
| |
| async def execute(self, request): |
| instance = self._client |
| if instance: |
| await instance.send(request) |
| return await instance.recv() |
| return None |
| |
| async def _start_client(self, url, max_retries=_CONNECT_MAX_RETRIES_DEFAULT, interval_between_retries=1): |
| if max_retries: |
| start = time.time() |
| try: |
| self._hooks.connecting(url) |
| connection = await websockets.connect(url, ping_timeout=_KEEP_ALIVE_TIMEOUT_IN_SECONDS, max_size=_MAX_MESSAGE_SIZE_IN_BYTES) |
| duration = round((time.time() - start) * 1000, 0) |
| self._hooks.success(duration) |
| return connection |
| except Exception: |
| duration = round((time.time() - start) * 1000, 0) |
| self._hooks.failure(duration) |
| self._hooks.retry(interval_between_retries) |
| time.sleep(interval_between_retries) |
| return await self._start_client(url, max_retries - 1, interval_between_retries + 1) |
| |
| self._hooks.abort(url) |
| raise Exception(f'Connecting to {url} failed.') |
| |
| async def _stop_client(self, instance): |
| if instance: |
| await instance.close() |
| |
| async def _start_server(self, command, url): |
| instance = None |
| if command: |
| start_time = time.time() |
| |
| command = ['stdbuf', '-o0', '-e0'] + command # disable buffering |
| instance = subprocess.Popen( |
| command, text=False, bufsize=0, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) |
| |
| # Loop to read the subprocess output with a timeout |
| lines = [] |
| while True: |
| if time.time() - start_time > _WEBSOCKET_SERVER_MESSAGE_TIMEOUT: |
| for line in lines: |
| print(line.decode('utf-8'), end='') |
| self._hooks.abort(url) |
| await self._stop_server(instance) |
| raise Exception( |
| f'Connecting to {url} failed. WebSocket startup has not been detected.') |
| |
| ready, _, _ = select.select([instance.stdout], [], [], 1) |
| if ready: |
| line = instance.stdout.readline() |
| if len(line): |
| lines.append(line) |
| if re.search(_WEBSOCKET_SERVER_MESSAGE, line.decode('utf-8')): |
| break # Exit the loop if the pattern is found |
| else: |
| continue |
| instance.stdout.close() |
| |
| return instance |
| |
| async def _stop_server(self, instance): |
| if instance: |
| instance.terminate() # sends SIGTERM |
| try: |
| instance.wait(_WEBSOCKET_SERVER_TERMINATE_TIMEOUT) |
| except subprocess.TimeoutExpired: |
| logging.debug( |
| 'Subprocess did not terminate on SIGTERM, killing it now') |
| instance.kill() |
| |
| def _make_server_connection_url(self, address: str, port: int): |
| return 'ws://' + address + ':' + str(port) |
| |
| def _make_server_startup_command(self, path: str, arguments: str, port: int): |
| if path is None: |
| return None |
| elif arguments is None: |
| return [path] + ['--port', str(port)] |
| else: |
| return [path] + [arg.strip() for arg in arguments.split(' ')] + ['--port', str(port)] |
