From e79dfabbea7ed1595a2f0f9157f7b434761eb27e Mon Sep 17 00:00:00 2001 From: Dariusz Suchojad Date: Mon, 27 Feb 2012 12:58:38 +0100 Subject: [PATCH] Made the WebSocketApp accept a mask_key function which is passed on down to the actual socket. --- test_websocket.py | 49 ++++++++++++++++++++++++++++++++++------------- websocket.py | 16 +++++++++++----- 2 files changed, 47 insertions(+), 18 deletions(-) diff --git a/test_websocket.py b/test_websocket.py index 36a3f6c..ab2b2f9 100644 --- a/test_websocket.py +++ b/test_websocket.py @@ -223,7 +223,7 @@ class WebSocketTest(unittest.TestCase): u = uuid.UUID(bytes=base64.b64decode(key)) self.assertEquals(4, u.version) -class WebSocketKeepRunningTest(unittest.TestCase): +class WebSocketAppTest(unittest.TestCase): class NotSetYet(object): """ A marker class for signalling that a value hasn't been set yet. @@ -232,13 +232,15 @@ class WebSocketKeepRunningTest(unittest.TestCase): def setUp(self): ws.enableTrace(TRACABLE) - WebSocketKeepRunningTest.keep_running_open = WebSocketKeepRunningTest.NotSetYet() - WebSocketKeepRunningTest.keep_running_close = WebSocketKeepRunningTest.NotSetYet() + WebSocketAppTest.keep_running_open = WebSocketAppTest.NotSetYet() + WebSocketAppTest.keep_running_close = WebSocketAppTest.NotSetYet() + WebSocketAppTest.get_mask_key_id = WebSocketAppTest.NotSetYet() def tearDown(self): - WebSocketKeepRunningTest.keep_running_open = WebSocketKeepRunningTest.NotSetYet() - WebSocketKeepRunningTest.keep_running_close = WebSocketKeepRunningTest.NotSetYet() + WebSocketAppTest.keep_running_open = WebSocketAppTest.NotSetYet() + WebSocketAppTest.keep_running_close = WebSocketAppTest.NotSetYet() + WebSocketAppTest.get_mask_key_id = WebSocketAppTest.NotSetYet() def testKeepRunning(self): """ A WebSocketApp should keep running as long as its self.keep_running @@ -249,25 +251,46 @@ class WebSocketKeepRunningTest(unittest.TestCase): """ Set the keep_running flag for later inspection and immediately close the connection. """ - WebSocketKeepRunningTest.keep_running_open = self.keep_running + WebSocketAppTest.keep_running_open = self.keep_running self.close() def on_close(self, *args, **kwargs): """ Set the keep_running flag for the test to use. """ - WebSocketKeepRunningTest.keep_running_close = self.keep_running + WebSocketAppTest.keep_running_close = self.keep_running app = ws.WebSocketApp('ws://echo.websocket.org/', on_open=on_open, on_close=on_close) app.run_forever() - self.assertFalse(isinstance(WebSocketKeepRunningTest.keep_running_open, - WebSocketKeepRunningTest.NotSetYet)) + self.assertFalse(isinstance(WebSocketAppTest.keep_running_open, + WebSocketAppTest.NotSetYet)) - self.assertFalse(isinstance(WebSocketKeepRunningTest.keep_running_close, - WebSocketKeepRunningTest.NotSetYet)) + self.assertFalse(isinstance(WebSocketAppTest.keep_running_close, + WebSocketAppTest.NotSetYet)) - self.assertEquals(True, WebSocketKeepRunningTest.keep_running_open) - self.assertEquals(False, WebSocketKeepRunningTest.keep_running_close) + self.assertEquals(True, WebSocketAppTest.keep_running_open) + self.assertEquals(False, WebSocketAppTest.keep_running_close) + + def testSockMaskKey(self): + """ A WebSocketApp should forward the received mask_key function down + to the actual socket. + """ + + def my_mask_key_func(): + pass + + def on_open(self, *args, **kwargs): + """ Set the value so the test can use it later on and immediately + close the connection. + """ + WebSocketAppTest.get_mask_key_id = id(self.get_mask_key) + self.close() + + app = ws.WebSocketApp('ws://echo.websocket.org/', on_open=on_open, get_mask_key=my_mask_key_func) + app.run_forever() + + # Note: We can't use 'is' for comparing the functions directly, need to use 'id'. + self.assertEquals(WebSocketAppTest.get_mask_key_id, id(my_mask_key_func)) if __name__ == "__main__": diff --git a/websocket.py b/websocket.py index ce9459d..6561f30 100644 --- a/websocket.py +++ b/websocket.py @@ -313,14 +313,17 @@ class WebSocket(object): >>> ws.recv() 'Hello, Server' >>> ws.close() + + get_mask_key: a callable to produce new mask keys, see the set_mask_key + function's docstring for more details """ - def __init__(self): + def __init__(self, get_mask_key = None): """ Initalize WebSocket object. """ self.connected = False self.io_sock = self.sock = socket.socket() - self.get_mask_key = None + self.get_mask_key = get_mask_key def set_mask_key(self, func): """ @@ -639,7 +642,7 @@ class WebSocketApp(object): """ def __init__(self, url, on_open = None, on_message = None, on_error = None, - on_close = None, keep_running = True): + on_close = None, keep_running = True, get_mask_key = None): """ url: websocket url. on_open: callable object which is called at opening websocket. @@ -656,14 +659,17 @@ class WebSocketApp(object): this function has one argument. The arugment is this class object. keep_running: a boolean flag indicating whether the app's main loop should keep running, defaults to True + get_mask_key: a callable to produce new mask keys, see the WebSocket.set_mask_key's + docstring for more information """ self.url = url self.on_open = on_open self.on_message = on_message self.on_error = on_error self.on_close = on_close - self.sock = None self.keep_running = keep_running + self.get_mask_key = get_mask_key + self.sock = None def send(self, data): """ @@ -686,7 +692,7 @@ class WebSocketApp(object): if self.sock: raise WebSocketException("socket is already opened") try: - self.sock = WebSocket() + self.sock = WebSocket(self.get_mask_key) self.sock.connect(self.url) self._run_with_no_err(self.on_open) while self.keep_running: