From b4d513066c6d7ca05f487b1c9e71a99d23a00c34 Mon Sep 17 00:00:00 2001
From: Dariusz Suchojad <dsuch-github@gefira.pl>
Date: Mon, 27 Feb 2012 12:28:52 +0100
Subject: [PATCH] Added a boolean keep_running flag to WebSocketApp so that the
 main loop can be interrupted.

---
 test_websocket.py | 46 ++++++++++++++++++++++++++++++++++++++++++++++
 websocket.py      |  8 ++++++--
 2 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/test_websocket.py b/test_websocket.py
index 3c68287..36a3f6c 100644
--- a/test_websocket.py
+++ b/test_websocket.py
@@ -222,6 +222,52 @@ class WebSocketTest(unittest.TestCase):
         key = ws._create_sec_websocket_key()
         u = uuid.UUID(bytes=base64.b64decode(key))
         self.assertEquals(4, u.version)
+
+class WebSocketKeepRunningTest(unittest.TestCase):
+
+    class NotSetYet(object):
+        """ A marker class for signalling that a value hasn't been set yet.
+        """
+    
+    def setUp(self):
+        ws.enableTrace(TRACABLE)
+        
+        WebSocketKeepRunningTest.keep_running_open = WebSocketKeepRunningTest.NotSetYet()
+        WebSocketKeepRunningTest.keep_running_close = WebSocketKeepRunningTest.NotSetYet()
+    
+    def tearDown(self):
+        
+        WebSocketKeepRunningTest.keep_running_open = WebSocketKeepRunningTest.NotSetYet()
+        WebSocketKeepRunningTest.keep_running_close = WebSocketKeepRunningTest.NotSetYet()
+        
+    def testKeepRunning(self):
+        """ A WebSocketApp should keep running as long as its self.keep_running
+        is not False (in the boolean context).
+        """
+
+        def on_open(self, *args, **kwargs):
+            """ Set the keep_running flag for later inspection and immediately
+            close the connection.
+            """
+            WebSocketKeepRunningTest.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
+        
+        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(WebSocketKeepRunningTest.keep_running_close, 
+                                    WebSocketKeepRunningTest.NotSetYet))
+        
+        self.assertEquals(True, WebSocketKeepRunningTest.keep_running_open)
+        self.assertEquals(False, WebSocketKeepRunningTest.keep_running_close)
         
 
 if __name__ == "__main__":
diff --git a/websocket.py b/websocket.py
index bcceedc..ce9459d 100644
--- a/websocket.py
+++ b/websocket.py
@@ -639,7 +639,7 @@ class WebSocketApp(object):
     """
     def __init__(self, url,
                  on_open = None, on_message = None, on_error = None, 
-                 on_close = None):
+                 on_close = None, keep_running = True):
         """
         url: websocket url.
         on_open: callable object which is called at opening websocket.
@@ -654,6 +654,8 @@ class WebSocketApp(object):
          The passing 2nd arugment is exception object.
        on_close: callable object which is called when closed the connection.
          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
         """
         self.url = url
         self.on_open = on_open
@@ -661,6 +663,7 @@ class WebSocketApp(object):
         self.on_error = on_error
         self.on_close = on_close
         self.sock = None
+        self.keep_running = keep_running
 
     def send(self, data):
         """
@@ -672,6 +675,7 @@ class WebSocketApp(object):
         """
         close websocket connection.
         """
+        self.keep_running = False
         self.sock.close()
 
     def run_forever(self):
@@ -685,7 +689,7 @@ class WebSocketApp(object):
             self.sock = WebSocket()
             self.sock.connect(self.url)
             self._run_with_no_err(self.on_open)
-            while True:
+            while self.keep_running:
                 data = self.sock.recv()
                 if data is None:
                     break