Merge "Raise proper exception at webscocket close" into stable/queens

tripleoclient/exceptions.py

@@ -28,6 +28,10 @@ class WebSocketTimeout(Exception):
     """Timed out waiting for messages on the websocket"""
 
 
+class WebSocketConnectionClosed(Exception):
+    """Websocket connection is closed before wait for messages"""
+
+
 class NotFound(Exception):
     """Resource not found"""
 

tripleoclient/plugin.py

@@ -144,20 +144,21 @@ class WebsocketClient(object):
         messages.
         """
 
-        if not self._ws.connected:
-            return
-
         if timeout is None:
             LOG.warning("Waiting for messages on queue '{}' with no timeout."
                         .format(self._queue_name))
 
-        self._ws.settimeout(timeout)
+        try:
+            self._ws.settimeout(timeout)
+            while True:
+                message = self.recv()
+                LOG.debug(message)
+                yield message['body']['payload']
 
-        while True:
-            try:
-                yield self.recv()['body']['payload']
-            except websocket.WebSocketTimeoutException:
-                raise exceptions.WebSocketTimeout()
+        except websocket.WebSocketTimeoutException:
+            raise exceptions.WebSocketTimeout()
+        except websocket.WebSocketConnectionClosedException:
+            raise exceptions.WebSocketConnectionClosed()
 
     def __enter__(self):
         """Return self to allow usage as a context manager"""

tripleoclient/tests/workflows/test_base.py

@@ -16,7 +16,7 @@ import mock
 
 from osc_lib.tests import utils
 
-from tripleoclient import exceptions
+from tripleoclient import exceptions as ex
 from tripleoclient.workflows import base
 
 
@@ -48,13 +48,28 @@ class TestBaseWorkflows(utils.TestCommand):
     def test_wait_for_messages_timeout(self):
         mistral = mock.Mock()
         websocket = mock.Mock()
-        websocket.wait_for_messages.side_effect = exceptions.WebSocketTimeout
+        websocket.wait_for_messages.side_effect = ex.WebSocketTimeout
         execution = mock.Mock()
         execution.id = 1
 
         messages = base.wait_for_messages(mistral, websocket, execution)
 
-        self.assertRaises(exceptions.WebSocketTimeout, list, messages)
+        self.assertRaises(ex.WebSocketTimeout, list, messages)
+
+        self.assertTrue(mistral.executions.get.called)
+        websocket.wait_for_messages.assert_called_with(timeout=None)
+
+    def test_wait_for_messages_connection_closed(self):
+        mistral = mock.Mock()
+        websocket = mock.Mock()
+        websocket.wait_for_messages.side_effect = ex.WebSocketConnectionClosed
+
+        execution = mock.Mock()
+        execution.id = 1
+
+        messages = base.wait_for_messages(mistral, websocket, execution)
+
+        self.assertRaises(ex.WebSocketConnectionClosed, list, messages)
 
         self.assertTrue(mistral.executions.get.called)
         websocket.wait_for_messages.assert_called_with(timeout=None)
@@ -78,5 +93,5 @@ class TestBaseWorkflows(utils.TestCommand):
         result.state = 'ERROR'
         mistral.action_executions.create = mock.Mock(return_value=result)
 
-        self.assertRaises(exceptions.WorkflowActionError,
+        self.assertRaises(ex.WorkflowActionError,
                           base.call_action, mistral, action)

tripleoclient/workflows/base.py

@@ -71,8 +71,8 @@ def wait_for_messages(mistral, websocket, execution, timeout=None):
             # Workflows should end with SUCCESS or ERROR statuses.
             if payload.get('status', 'RUNNING') != "RUNNING" or \
                     mistral.executions.get(execution.id).state != "RUNNING":
-                raise StopIteration
-    except exceptions.WebSocketTimeout:
+                return
+    except (exceptions.WebSocketTimeout, exceptions.WebSocketConnectionClosed):
         check_execution_status(mistral, execution.id)
         raise