From 0a5ce3641b8d3cfea847667b004c0ee5ad2cf52b Mon Sep 17 00:00:00 2001 From: Denis Soldatov Date: Thu, 13 Feb 2020 21:28:19 +0300 Subject: [PATCH 1/3] add init-message param --- wsstat/clients.py | 4 ++++ wsstat/main.py | 9 +++++++++ 2 files changed, 13 insertions(+) diff --git a/wsstat/clients.py b/wsstat/clients.py index f09aec9..7ac0ec9 100644 --- a/wsstat/clients.py +++ b/wsstat/clients.py @@ -68,6 +68,7 @@ def __init__(self, websocket_url, **kwargs): self.total_connections = kwargs.get('total_connections', 250) self._exiting = False self.extra_headers = None + self.init_message = kwargs.get('init_message') # Asyncio stuff self.loop = asyncio.get_event_loop() @@ -140,6 +141,9 @@ def create_websocket_connection(self): # Await the connection to complete successfully websocket = yield from websockets.connect(**connection_args) websocket.connection_time = time.time() - start_time + + if self.init_message: + yield from websocket.send(self.init_message) break except BaseException as e: diff --git a/wsstat/main.py b/wsstat/main.py index 6f67dfa..0b6b391 100644 --- a/wsstat/main.py +++ b/wsstat/main.py @@ -52,6 +52,15 @@ def parse_args(): action="store_true" ) + parser.add_argument( + '-m', "--init-message", + help="Initial message that a client will send to the websocket server", + dest="init_message", + action="store", + default=None, + type=str + ) + args = parser.parse_args() if not args.websocket_url and not args.demo: From f70b482fcb0c9f1a986688e62b3d88efea1a69f9 Mon Sep 17 00:00:00 2001 From: Denis Soldatov Date: Thu, 13 Feb 2020 21:32:09 +0300 Subject: [PATCH 2/3] update description --- README.md | 5 ++++- wsstat/main.py | 2 +- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 59f815e..3cb41b1 100644 --- a/README.md +++ b/README.md @@ -54,7 +54,7 @@ Other than that, you can adjust the total number of connected clients with `-n`, ``` $ wsstat -h usage: wsstat [-h] [-n TOTAL_CONNECTIONS] [-c MAX_CONNECTING_SOCKETS] - [-H HEADER] [--demo] [-i] + [-H HEADER] [--demo] [-i] [-m INIT_MESSAGE] [websocket_url] positional arguments: @@ -71,4 +71,7 @@ optional arguments: Pass a custom header with each websocket connection --demo Start a demo websocket server and point wsstat at it -i, --insecure Don't validate SSL certificates on websocket servers + -m INIT_MESSAGE, --init-message INIT_MESSAGE + Initial message that a client will send to the + websocket server after connection ``` diff --git a/wsstat/main.py b/wsstat/main.py index 0b6b391..5b00ab5 100644 --- a/wsstat/main.py +++ b/wsstat/main.py @@ -54,7 +54,7 @@ def parse_args(): parser.add_argument( '-m', "--init-message", - help="Initial message that a client will send to the websocket server", + help="Initial message that a client will send to the websocket server after connection", dest="init_message", action="store", default=None, From 83dd683ba42a1bc62606fda875514a716443f544 Mon Sep 17 00:00:00 2001 From: Denis Soldatov Date: Fri, 14 Feb 2020 14:10:17 +0300 Subject: [PATCH 3/3] fix test --- test/test_coros.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/test_coros.py b/test/test_coros.py index 4a7ea94..c39717d 100644 --- a/test/test_coros.py +++ b/test/test_coros.py @@ -91,7 +91,7 @@ def test_valid_connection(self): all_messages = self.arg_list_to_string(self.client.logger.log.call_args_list) assert "Connecting" in all_messages assert "Connected" in all_messages - assert "connection is closed" in all_messages + assert "Peace out homie!" in all_messages assert self.identifier in all_messages connected_websocket = self.client.sockets[self.identifier]