From 41a2687b9be14165b31dec1e8a0cb067bb202b8e Mon Sep 17 00:00:00 2001
From: vrde <agranzot@gmail.com>
Date: Mon, 27 Aug 2018 15:00:02 +0200
Subject: [PATCH] Problem: websocket fails with ELECTION transaction (#2482)

Solution: have a more general approach to process transaction types. If
a transaction does *not* contain `asset.id`, then the `id` of the
`asset` is the `id` of the transaction.
---
 bigchaindb/web/websocket_server.py | 24 ++++++++++++++----------
 tests/web/test_websocket_server.py | 27 +++++++++++++++++++++++++++
 2 files changed, 41 insertions(+), 10 deletions(-)

diff --git a/bigchaindb/web/websocket_server.py b/bigchaindb/web/websocket_server.py
index 8aea0e20..f0a9f886 100644
--- a/bigchaindb/web/websocket_server.py
+++ b/bigchaindb/web/websocket_server.py
@@ -48,6 +48,17 @@ def _multiprocessing_to_asyncio(in_queue, out_queue, loop):
         loop.call_soon_threadsafe(out_queue.put_nowait, value)
 
 
+def eventify_block(block):
+    for tx in block['transactions']:
+        try:
+            asset_id = tx['asset']['id']
+        except KeyError:
+            asset_id = tx['id']
+        yield {'height': block['height'],
+               'asset_id': asset_id,
+               'transaction_id': tx['id']}
+
+
 class Dispatcher:
     """Dispatch events to websockets.
 
@@ -99,17 +110,10 @@ class Dispatcher:
                 str_buffer.append(event)
 
             elif event.type == EventTypes.BLOCK_VALID:
-                block = event.data
+                str_buffer = map(json.dumps, eventify_block(event.data))
 
-                for tx in block['transactions']:
-                    asset_id = tx['id'] if tx['operation'] == 'CREATE' else tx['asset']['id']
-                    data = {'height': block['height'],
-                            'asset_id': asset_id,
-                            'transaction_id': tx['id']}
-                    str_buffer.append(json.dumps(data))
-
-            for _, websocket in self.subscribers.items():
-                for str_item in str_buffer:
+            for str_item in str_buffer:
+                for _, websocket in self.subscribers.items():
                     yield from websocket.send_str(str_item)
 
 
diff --git a/tests/web/test_websocket_server.py b/tests/web/test_websocket_server.py
index 1a2b55af..a3de508a 100644
--- a/tests/web/test_websocket_server.py
+++ b/tests/web/test_websocket_server.py
@@ -21,6 +21,33 @@ class MockWebSocket:
         self.received.append(s)
 
 
+def test_eventify_block_works_with_any_transaction():
+    from bigchaindb.web.websocket_server import eventify_block
+
+    block = {
+        'height': 1,
+        'transactions': [{
+            'id': 1
+        }, {
+            'id': 2,
+            'asset': {'id': 1}
+            }]
+        }
+
+    expected_events = [{
+            'height': 1,
+            'asset_id': 1,
+            'transaction_id': 1
+        }, {
+            'height': 1,
+            'asset_id': 1,
+            'transaction_id': 2
+        }]
+
+    for event, expected in zip(eventify_block(block), expected_events):
+        assert event == expected
+
+
 @asyncio.coroutine
 def test_bridge_sync_async_queue(loop):
     from bigchaindb.web.websocket_server import _multiprocessing_to_asyncio