diff --git a/test/websocket_server.py b/test/websocket_server.py
index 7b461b00bd..99a23e4e0a 100755
--- a/test/websocket_server.py
+++ b/test/websocket_server.py
@@ -40,7 +40,7 @@ async def handle(websocket, path):
         print(e)
 
 
-if __name__ == '__main__':
+async def main() -> None:
     port = int(sys.argv[1])
     use_ssl = sys.argv[2] != '0'
     min_interval = sys.argv[3]
@@ -52,6 +52,9 @@ async def handle(websocket, path):
     else:
         ssl_context = None
 
-    start_server = websockets.serve(handle, '127.0.0.1', port, ssl=ssl_context)
-    asyncio.get_event_loop().run_until_complete(start_server)
+    await websockets.serve(handle, '127.0.0.1', port, ssl=ssl_context)
+
+if __name__ == '__main__':
+    loop = asyncio.get_event_loop()
+    asyncio.get_event_loop().run_until_complete(main())
     asyncio.get_event_loop().run_forever()