Support init message

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
 ```

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]

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:

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: