Update websocket_test.py