diff --git a/tests/load/locustfiles/stored.py b/tests/load/locustfiles/stored.py index 07170d724..5a378c5d1 100644 --- a/tests/load/locustfiles/stored.py +++ b/tests/load/locustfiles/stored.py @@ -16,11 +16,9 @@ from logging import Logger from typing import Any, TypeAlias -import gevent import websocket from args import parse_wait_time from exceptions import ZeroStatusRequestError -from gevent import Greenlet from locust import FastHttpUser, events, task from locust.exception import LocustError from models import ( @@ -38,9 +36,6 @@ Message: TypeAlias = HelloMessage | NotificationMessage | RegisterMessage | UnregisterMessage Record: TypeAlias = HelloRecord | NotificationRecord | RegisterRecord -# Set to 'True' to view the verbose connection information for the web socket -#websocket.enableTrace(False) -#websocket.setdefaulttimeout(5) logger: Logger = logging.getLogger("StoredNotifAutopushUser") @@ -65,8 +60,7 @@ def __init__(self, environment) -> None: self.register_records: list[RegisterRecord] = [] self.unregister_records: list[RegisterRecord] = [] self.uaid: str = "" - self.ws: WebSocketApp | None = None - self.ws_greenlet: Greenlet | None = None + self.ws: WebSocket | None = None def wait_time(self): return self.environment.autopush_wait_time(self) @@ -82,8 +76,6 @@ def on_stop(self) -> Any: self.send_unregister(self.ws, channel_id) self.ws.close() self.ws = None - if self.ws_greenlet: - gevent.kill(self.ws_greenlet) def on_ws_open(self, ws: WebSocket) -> None: """Called when opening a WebSocket. @@ -362,7 +354,7 @@ def send_register(self, ws: WebSocket, channel_id: str) -> None: self.register_records.append(record) self.send(ws, message_type, data) - def send_unregister(self, ws: WebSocketApp, channel_id: str) -> None: + def send_unregister(self, ws: WebSocket, channel_id: str) -> None: """Send an 'unregister' message to Autopush. Args: