Refactor the websocket test