diff --git a/tests/load/locustfiles/locustfile.py b/tests/load/locustfiles/locustfile.py index f79079a98..3b976ec90 100644 --- a/tests/load/locustfiles/locustfile.py +++ b/tests/load/locustfiles/locustfile.py @@ -28,11 +28,12 @@ NotificationRecord, RegisterMessage, RegisterRecord, + UnregisterMessage, ) from pydantic import ValidationError from websocket import WebSocketApp, WebSocketConnectionClosedException -Message: TypeAlias = HelloMessage | NotificationMessage | RegisterMessage +Message: TypeAlias = HelloMessage | NotificationMessage | RegisterMessage | UnregisterMessage Record: TypeAlias = HelloRecord | NotificationRecord | RegisterRecord # Set to 'True' to view the verbose connection information for the web socket @@ -85,6 +86,7 @@ def __init__(self, environment) -> None: self.hello_record: HelloRecord | None = None self.notification_records: list[NotificationRecord] = [] self.register_records: list[RegisterRecord] = [] + self.unregister_records: list[RegisterRecord] = [] self.uaid: str = "" self.ws: WebSocketApp | None = None self.ws_greenlet: Greenlet | None = None @@ -99,6 +101,8 @@ def on_start(self) -> Any: def on_stop(self) -> Any: """Called when a User stops running.""" if self.ws: + for channel_id in self.channels.keys(): + self.send_unregister(self.ws, channel_id) self.ws.close() self.ws = None if self.ws_greenlet: @@ -122,14 +126,12 @@ def on_ws_message(self, ws: WebSocketApp, data: str) -> None: message: Message | None = self.recv(data) if isinstance(message, HelloMessage): self.uaid = message.uaid - # In order to assure notifications are sent on startup, users are - # required to be subscribed to a channel - if not self.channels: - self.send_register(ws) elif isinstance(message, NotificationMessage): self.send_ack(ws, message.channelID, message.version) elif isinstance(message, RegisterMessage): self.channels[message.channelID] = message.pushEndpoint + elif isinstance(message, UnregisterMessage): + del self.channels[message.channelID] def on_ws_error(self, ws: WebSocketApp, error: Exception) -> None: """Called when there is a WebSocketApp error or if an exception is raised in a WebSocket @@ -172,6 +174,26 @@ def send_notification(self): endpoint_url: str = random.choice(list(self.channels.values())) self.post_notification(endpoint_url) + @task(weight=1) + def subscribe(self): + """Subscribes a user to an Autopush channel.""" + if not self.ws: + logger.debug("Task 'subscribe' skipped.") + return + + channel_id: str = str(uuid.uuid4()) + self.send_register(self.ws, channel_id) + + @task(weight=1) + def unsubscribe(self): + """Unsubscribes a user from an Autopush channel.""" + if not self.ws or not self.channels: + logger.debug("Task 'unsubscribe' skipped.") + return + + channel_id: str = random.choice(list(self.channels.keys())) + self.send_unregister(self.ws, channel_id) + def connect(self) -> None: """Creates the WebSocketApp that will run indefinitely.""" self.ws = websocket.WebSocketApp( @@ -250,9 +272,16 @@ def recv(self, data: str) -> Message | None: ) case "register": message = RegisterMessage(**message_dict) - message_channel_id: str = message.channelID # type: ignore[union-attr] + register_chid: str = message.channelID # type: ignore[union-attr] + record = next( + (r for r in self.register_records if r.channel_id == register_chid), + None, + ) + case "unregister": + message = UnregisterMessage(**message_dict) + unregister_chid: str = message.channelID # type: ignore[union-attr] record = next( - (r for r in self.register_records if r.channel_id == message_channel_id), + (r for r in self.unregister_records if r.channel_id == unregister_chid), None, ) case _: @@ -318,27 +347,42 @@ def send_hello(self, ws: WebSocketApp) -> None: self.hello_record = HelloRecord(send_time=time.perf_counter()) self.send(ws, message_type, data) - def send_register(self, ws: WebSocketApp) -> None: + def send_register(self, ws: WebSocketApp, channel_id: str) -> None: """Send a 'register' message to Autopush. Args: ws: WebSocket class object + channel_id: Notification message channel ID Raises: WebSocketException: Error raised by the WebSocket client """ message_type: str = "register" - channel_id: str = str(uuid.uuid4()) data: dict[str, Any] = dict(messageType=message_type, channelID=channel_id) record = RegisterRecord(send_time=time.perf_counter(), channel_id=channel_id) self.register_records.append(record) self.send(ws, message_type, data) + def send_unregister(self, ws: WebSocketApp, channel_id: str) -> None: + """Send an 'unregister' message to Autopush. + + Args: + ws: WebSocket class object + channel_id: Notification message channel ID + Raises: + WebSocketException: Error raised by the WebSocket client + """ + message_type: str = "unregister" + data: dict[str, Any] = dict(messageType=message_type, channelID=channel_id) + record = RegisterRecord(send_time=time.perf_counter(), channel_id=channel_id) + self.unregister_records.append(record) + self.send(ws, message_type, data) + def send(self, ws: WebSocketApp, message_type: str, data: dict[str, Any]) -> None: """Send a message to Autopush. Args: ws: WebSocket class object - message_type: Message type. Examples: 'ack', 'hello' or 'register' + message_type: Message type. Examples: 'ack', 'hello', 'register' or 'unregister' data: Message data Raises: WebSocketException: Error raised by the WebSocket client diff --git a/tests/load/locustfiles/models.py b/tests/load/locustfiles/models.py index 856d8d806..5b36d98f2 100644 --- a/tests/load/locustfiles/models.py +++ b/tests/load/locustfiles/models.py @@ -56,3 +56,11 @@ class RegisterRecord(BaseModel): send_time: float channel_id: str + + +class UnregisterMessage(BaseModel): + """Autopush 'unregister' response message.""" + + messageType: Literal["unregister"] + channelID: str + status: Literal[200]