From 6f94cfd791e434e9932da18d785feb934d802a44 Mon Sep 17 00:00:00 2001
From: Thomas Herve <therve@redhat.com>
Date: Thu, 3 Nov 2016 09:57:58 +0100
Subject: [PATCH] Add convenience function for events over websocket

This adds an alternative method of poll_for_events using a websocket
connection.

Change-Id: I1c46af3ab98b6d108af976e4fab5282c8756a626
---
 heatclient/common/event_utils.py          | 25 ++++++++++++++
 heatclient/common/utils.py                |  6 ++--
 heatclient/tests/unit/test_event_utils.py | 42 +++++++++++++++++++++++
 3 files changed, 71 insertions(+), 2 deletions(-)

diff --git a/heatclient/common/event_utils.py b/heatclient/common/event_utils.py
index 9208d0c3..0e93bcee 100644
--- a/heatclient/common/event_utils.py
+++ b/heatclient/common/event_utils.py
@@ -18,6 +18,7 @@ import time
 from heatclient._i18n import _
 from heatclient.common import utils
 import heatclient.exc as exc
+from heatclient.v1 import events as events_mod
 
 
 def get_hook_events(hc, stack_id, event_args, nested_depth=0,
@@ -216,3 +217,27 @@ def poll_for_events(hc, stack_name, action=None, poll_period=5, marker=None,
             no_event_polls = 0
 
         time.sleep(poll_period)
+
+
+def wait_for_events(ws, stack_name, out=None):
+    """Receive events over the passed websocket and wait for final status."""
+    msg_template = _("\n Stack %(name)s %(status)s \n")
+    if not out:
+        out = sys.stdout
+    event_log_context = utils.EventLogContext()
+    while True:
+        data = ws.recv()['body']
+        event = events_mod.Event(None, data['payload'], True)
+        # Keep compatibility with the HTTP API
+        event.event_time = data['timestamp']
+        event.resource_status = '%s_%s' % (event.resource_action,
+                                           event.resource_status)
+        events_log = utils.event_log_formatter([event], event_log_context)
+        out.write(events_log)
+        out.write('\n')
+        if data['payload']['resource_name'] == stack_name:
+            stack_status = data['payload']['resource_status']
+            if stack_status in ('COMPLETE', 'FAILED'):
+                msg = msg_template % dict(
+                    name=stack_name, status=event.resource_status)
+                return '%s_%s' % (event.resource_action, stack_status), msg
diff --git a/heatclient/common/utils.py b/heatclient/common/utils.py
index 91410acd..3ebfede1 100644
--- a/heatclient/common/utils.py
+++ b/heatclient/common/utils.py
@@ -125,6 +125,8 @@ class EventLogContext(object):
         # future calls to build_resource_name
 
         def get_stack_id():
+            if getattr(event, 'stack_id', None) is not None:
+                return event.stack_id
             for l in getattr(event, 'links', []):
                 if l.get('rel') == 'stack':
                     if 'href' not in l:
@@ -135,8 +137,8 @@ class EventLogContext(object):
         stack_id = get_stack_id()
         if not stack_id:
             return res_name
-        phys_id = getattr(event, 'physical_resource_id')
-        status = getattr(event, 'resource_status')
+        phys_id = getattr(event, 'physical_resource_id', None)
+        status = getattr(event, 'resource_status', None)
 
         is_stack_event = stack_id == phys_id
         if is_stack_event:
diff --git a/heatclient/tests/unit/test_event_utils.py b/heatclient/tests/unit/test_event_utils.py
index 32e44737..58f2e288 100644
--- a/heatclient/tests/unit/test_event_utils.py
+++ b/heatclient/tests/unit/test_event_utils.py
@@ -18,6 +18,15 @@ from heatclient.v1 import events as hc_ev
 from heatclient.v1 import resources as hc_res
 
 
+class FakeWebSocket(object):
+
+    def __init__(self, events):
+        self.events = events
+
+    def recv(self):
+        return self.events.pop(0)
+
+
 class ShellTestEventUtils(testtools.TestCase):
     @staticmethod
     def _mock_resource(resource_id, nested_id=None):
@@ -245,3 +254,36 @@ class ShellTestEventUtils(testtools.TestCase):
             mock_client, 'astack', action='CREATE', poll_period=0)
         self.assertEqual('CREATE_FAILED', stack_status)
         self.assertEqual('\n Stack astack CREATE_FAILED \n', msg)
+
+    def test_wait_for_events(self):
+        ws = FakeWebSocket([
+            {'body': {
+                'timestamp': '2014-01-06T16:14:26Z',
+                'payload': {'resource_action': 'CREATE',
+                            'resource_status': 'COMPLETE',
+                            'resource_name': 'mystack',
+                            'physical_resource_id': 'stackid1',
+                            'stack_id': 'stackid1'}}}])
+        stack_status, msg = event_utils.wait_for_events(ws, 'mystack')
+        self.assertEqual('CREATE_COMPLETE', stack_status)
+        self.assertEqual('\n Stack mystack CREATE_COMPLETE \n', msg)
+
+    def test_wait_for_events_failed(self):
+        ws = FakeWebSocket([
+            {'body': {
+                'timestamp': '2014-01-06T16:14:23Z',
+                'payload': {'resource_action': 'CREATE',
+                            'resource_status': 'IN_PROGRESS',
+                            'resource_name': 'mystack',
+                            'physical_resource_id': 'stackid1',
+                            'stack_id': 'stackid1'}}},
+            {'body': {
+                'timestamp': '2014-01-06T16:14:26Z',
+                'payload': {'resource_action': 'CREATE',
+                            'resource_status': 'FAILED',
+                            'resource_name': 'mystack',
+                            'physical_resource_id': 'stackid1',
+                            'stack_id': 'stackid1'}}}])
+        stack_status, msg = event_utils.wait_for_events(ws, 'mystack')
+        self.assertEqual('CREATE_FAILED', stack_status)
+        self.assertEqual('\n Stack mystack CREATE_FAILED \n', msg)