From 531c1504ab6d43066d1c17a1fb87b5a060d7d932 Mon Sep 17 00:00:00 2001
From: Quique Llorente <ellorent@redhat.com>
Date: Wed, 26 Sep 2018 12:09:41 +0200
Subject: [PATCH] Raise proper exception at webscocket close

After fixing a bug with https://review.openstack.org/#/c/603802/ the
return introduced there is bogus, we have to raise a proper exception
and handle it like timeouts, to get all the mistral allright.

Change-Id: Idcdbd38129f5694c5452f3f8aca0388df80476b2
(cherry picked from commit fee9f8cf15d708e4b3b8bdf0f0ddd6df8ce5c223)
---
 tripleoclient/exceptions.py                |  4 ++++
 tripleoclient/plugin.py                    | 19 +++++++++---------
 tripleoclient/tests/workflows/test_base.py | 23 ++++++++++++++++++----
 tripleoclient/workflows/base.py            |  4 ++--
 4 files changed, 35 insertions(+), 15 deletions(-)

diff --git a/tripleoclient/exceptions.py b/tripleoclient/exceptions.py
index adafe7905..5bcca1159 100644
--- a/tripleoclient/exceptions.py
+++ b/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"""
 
diff --git a/tripleoclient/plugin.py b/tripleoclient/plugin.py
index e467233f4..bdfaaebaa 100644
--- a/tripleoclient/plugin.py
+++ b/tripleoclient/plugin.py
@@ -146,20 +146,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"""
diff --git a/tripleoclient/tests/workflows/test_base.py b/tripleoclient/tests/workflows/test_base.py
index c1248b313..aec21ac49 100644
--- a/tripleoclient/tests/workflows/test_base.py
+++ b/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)
diff --git a/tripleoclient/workflows/base.py b/tripleoclient/workflows/base.py
index 3a1514f94..f3a784061 100644
--- a/tripleoclient/workflows/base.py
+++ b/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