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/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] 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..5b00ab5 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 after connection", + dest="init_message", + action="store", + default=None, + type=str + ) + args = parser.parse_args() if not args.websocket_url and not args.demo: