So I'm doing a unittest of a tornado server I programmed which has a websocketHandler
. This is my testing code:
def test_unauthorized_websocket(self):
message = "Hey bro"
ws = websocket.create_connection('ws://localhost:8888/ws', header=['Authorization: false'])
send_dict = {'command': test_command, 'message': message}
serialized_dict = json.dumps(send_dict)
ws.send(serialized_dict)
response = ws.recv()
#test response with assert
My goal with this test was to prove that my tornado server correctly refuses and closes this websocket connection because of wrong authentication header.
This is my tornado websocketHandler
code:
class WebSocketHandler(tornado.websocket.WebSocketHandler):
def open(self):
#some code
headers = self.request.headers
try:
auth = headers['Authorization']
except KeyError:
self.close(code = 1002, reason = "Unauthorized websocket")
print("IT'S CLOSED")
return
if auth == "true":
print("Authorized Websocket!")
#some code
else:
print("Unauthorized Websocket... :-(")
self.close(code = 1002, reason = "Unauthorized websocket")
So when the authentication is wrong, self.close()
is called (not sure I need code and reason). This should close the websocket. But this doesn't actually happens in the "client" side. After I call create_connection()
in the "client" the ws.connected
variable is still True, and when I do ws.send()
the on_message
method of the websocketHandler
is still called and when it tries to make a response with self.write_message()
it raises a WebSocketClosedError
. And only then the "client" actually closes its side of the websocket, ws.recv()
returns nothing and ws.connected
turns to False after that.
Is there a way I can communicate to the client side (through handshake headers or something) that the websocket is meant to be closed earlier on its side?
You can override prepare()
and raise a tornado.web.HTTPError
, instead of overriding open()
and calling self.close()
. This will reject the connection at the first opportunity and this will be reported on the client side as a part of create_connection()
instead of on a later read.