Restore IoTHub metadata in EventHub (#42)

Hazhzeng · web-flow · commit 31927c13ec59 · 2020-02-10T10:21:57.000-08:00
* Add iothub metadata in single &amp; multiple

* Fix format

* Add unit tests

* Fix linting

* Trim off unnecessary changes

* Change cardinary to cardinality
diff --git a/.flake8 b/.flake8
@@ -2,4 +2,4 @@
 ignore = W503,E402,E731
 exclude =
     .git, __pycache__, build, dist, .eggs, .github, .local,
-    Samples, azure/functions/_thirdparty, docs/
+    Samples, azure/functions/_thirdparty, docs/, .venv*/, .env*/, .vscode/
diff --git a/azure/functions/eventhub.py b/azure/functions/eventhub.py
@@ -105,12 +105,10 @@ def decode(
                       typing.List[_eventhub.EventHubEvent]]:
         data_type = data.type
 
-        if (data_type == 'string' or data_type == 'bytes'
-                or data_type == 'json'):
+        if cls._is_cardinality_one(trigger_metadata):
             return cls.decode_single_event(data, trigger_metadata)
 
-        elif (data_type == 'collection_bytes'
-                or data_type == 'collection_string'):
+        elif cls._is_cardinality_many(trigger_metadata):
             return cls.decode_multiple_events(data, trigger_metadata)
 
         else:
@@ -129,13 +127,6 @@ def decode_single_event(cls, data,
         elif data.type == 'json':
             body = data.value.encode('utf-8')
 
-        iothub_metadata = {}
-        for f in trigger_metadata:
-            if f.startswith('iothub-'):
-                v = cls._decode_trigger_metadata_field(
-                    trigger_metadata, f, python_type=str)
-                iothub_metadata[f[len('iothub-'):]] = v
-
         return _eventhub.EventHubEvent(
             body=body,
             enqueued_time=cls._parse_datetime_metadata(
@@ -146,7 +137,7 @@ def decode_single_event(cls, data,
                 trigger_metadata, 'SequenceNumber', python_type=int),
             offset=cls._decode_trigger_metadata_field(
                 trigger_metadata, 'Offset', python_type=str),
-            iothub_metadata=iothub_metadata
+            iothub_metadata=cls._decode_iothub_metadata(trigger_metadata)
         )
 
     @classmethod
@@ -159,6 +150,10 @@ def decode_multiple_events(
         elif data.type == 'collection_string':
             parsed_data = data.value.string
 
+        # Input Trigger IotHub Event
+        elif data.type == 'json':
+            parsed_data = json.loads(data.value)
+
         sys_props = trigger_metadata.get('SystemPropertiesArray')
 
         parsed_sys_props = json.loads(sys_props.value)
@@ -180,17 +175,79 @@ def decode_multiple_events(
                 expected_type=int)
 
             event = _eventhub.EventHubEvent(
-                body=parsed_data[i],
+                body=cls._marshall_event_body(parsed_data[i], data.type),
                 enqueued_time=cls._parse_datetime(enqueued_time),
                 partition_key=cls._decode_typed_data(
                     partition_key, python_type=str),
                 sequence_number=cls._decode_typed_data(
                     sequence_number, python_type=int),
                 offset=cls._decode_typed_data(
                     offset, python_type=int),
-                iothub_metadata={}
+                iothub_metadata=cls._extract_iothub_from_dict(
+                    parsed_sys_props[i])
             )
 
             events.append(event)
 
         return events
+
+    @classmethod
+    def _marshall_event_body(self, parsed_data, data_type):
+        # In IoTHub, when setting the eventhub using cardinality = 'many'
+        # The data is wrapped inside a json (e.g. '[{ "device-id": "1" }]')
+
+        # Previously, since the IoTHub events has a 'json' datatype,
+        # it is handled as single_event by mistake and our users handle the
+        # data parsing. And we want to keep the same behavior here.
+        if data_type == 'json':
+            return json.dumps(parsed_data).encode('utf-8')
+
+        return parsed_data
+
+    @classmethod
+    def _decode_iothub_metadata(
+            cls, trigger_metadata) -> typing.Dict[str, str]:
+        # Try extracting iothub_metadata from trigger_metadata
+        iothub_metadata = cls._extract_iothub_from_trigger_metadata(
+            trigger_metadata)
+
+        # Try extracting iothub_metadata from SystemProperties
+        if not iothub_metadata and trigger_metadata.get('SystemProperties'):
+            iothub_metadata = cls._extract_iothub_from_system_properties(
+                trigger_metadata['SystemProperties'].value)
+
+        return iothub_metadata
+
+    @classmethod
+    def _extract_iothub_from_trigger_metadata(
+            cls, metadict: typing.Dict[str, str]) -> typing.Dict[str, str]:
+        iothub_metadata = {}
+        for f in metadict:
+            if f.startswith('iothub-'):
+                v = cls._decode_trigger_metadata_field(
+                    metadict, f, python_type=str)
+                iothub_metadata[f[len('iothub-'):]] = v
+        return iothub_metadata
+
+    @classmethod
+    def _extract_iothub_from_system_properties(
+            cls, system_properties_string: str) -> typing.Dict[str, str]:
+        system_properties = json.loads(system_properties_string)
+        return cls._extract_iothub_from_dict(system_properties)
+
+    @classmethod
+    def _extract_iothub_from_dict(
+            cls, metadict: typing.Dict[str, str]) -> typing.Dict[str, str]:
+        iothub_metadata = {}
+        for f in metadict:
+            if f.startswith('iothub-'):
+                iothub_metadata[f[len('iothub-'):]] = metadict[f]
+        return iothub_metadata
+
+    @classmethod
+    def _is_cardinality_many(cls, trigger_metadata) -> bool:
+        return 'SystemPropertiesArray' in trigger_metadata
+
+    @classmethod
+    def _is_cardinality_one(cls, trigger_metadata) -> bool:
+        return 'SystemProperties' in trigger_metadata
diff --git a/tests/test_eventhub.py b/tests/test_eventhub.py
@@ -1,11 +1,27 @@
 from typing import List
 import unittest
+import json
+from unittest.mock import patch
+from datetime import datetime
 
 import azure.functions as func
 import azure.functions.eventhub as azf_eh
+import azure.functions.meta as meta
+
+
+class CollectionBytes:
+    def __init__(self, data: List[bytes]):
+        self.bytes = data
+
+
+class CollectionString:
+    def __init__(self, data: List[str]):
+        self.string = list(map(lambda x: x.encode('utf-8'), data))
 
 
 class TestEventHub(unittest.TestCase):
+    MOCKED_ENQUEUE_TIME = datetime.utcnow()
+
     def test_eventhub_input_type(self):
         check_input_type = (
             azf_eh.EventHubConverter.check_input_type_annotation
@@ -26,3 +42,187 @@ def test_eventhub_output_type(self):
         self.assertFalse(check_output_type(func.EventHubEvent))
         self.assertFalse(check_output_type(List[bytes]))
         self.assertFalse(check_output_type(List[func.EventHubEvent]))
+
+    @patch('azure.functions.eventhub.EventHubTriggerConverter'
+           '.decode_single_event')
+    @patch('azure.functions.eventhub.EventHubTriggerConverter'
+           '.decode_multiple_events')
+    def test_eventhub_decode_single_event(self, dme_mock, dse_mock):
+        azf_eh.EventHubTriggerConverter.decode(
+            data=self._generate_single_iothub_datum(),
+            trigger_metadata=self._generate_single_trigger_metadatum()
+        )
+        dse_mock.assert_called_once()
+        dme_mock.assert_not_called()
+
+    @patch('azure.functions.eventhub.EventHubTriggerConverter'
+           '.decode_single_event')
+    @patch('azure.functions.eventhub.EventHubTriggerConverter'
+           '.decode_multiple_events')
+    def test_eventhub_decode_multiple_events(self, dme_mock, dse_mock):
+        azf_eh.EventHubTriggerConverter.decode(
+            data=self._generate_multiple_iothub_data(),
+            trigger_metadata=self._generate_multiple_trigger_metadata()
+        )
+        dse_mock.assert_not_called()
+        dme_mock.assert_called_once()
+
+    def test_eventhub_trigger_single_event_json(self):
+        result = azf_eh.EventHubTriggerConverter.decode(
+            data=self._generate_single_iothub_datum('json'),
+            trigger_metadata=self._generate_single_trigger_metadatum()
+        )
+        # Result body always has the datatype of bytes
+        self.assertEqual(
+            result.get_body().decode('utf-8'), '{"device-status": "good"}'
+        )
+        self.assertEqual(result.enqueued_time, self.MOCKED_ENQUEUE_TIME)
+
+    def test_eventhub_trigger_single_event_bytes(self):
+        result = azf_eh.EventHubTriggerConverter.decode(
+            data=self._generate_single_iothub_datum('bytes'),
+            trigger_metadata=self._generate_single_trigger_metadatum()
+        )
+        self.assertEqual(
+            result.get_body().decode('utf-8'), '{"device-status": "good"}'
+        )
+        self.assertEqual(result.enqueued_time, self.MOCKED_ENQUEUE_TIME)
+
+    def test_iothub_metadata_single_event(self):
+        result = azf_eh.EventHubTriggerConverter.decode(
+            data=self._generate_single_iothub_datum('json'),
+            trigger_metadata=self._generate_single_trigger_metadatum()
+        )
+        self.assertIsNotNone(result.iothub_metadata)
+        self.assertEqual(
+            result.iothub_metadata['connection-device-id'], 'MyTestDevice'
+        )
+
+    def test_eventhub_trigger_single_event_string(self):
+        result = azf_eh.EventHubTriggerConverter.decode(
+            data=self._generate_single_iothub_datum('string'),
+            trigger_metadata=self._generate_single_trigger_metadatum()
+        )
+        self.assertEqual(
+            result.get_body().decode('utf-8'), '{"device-status": "good"}'
+        )
+        self.assertEqual(result.enqueued_time, self.MOCKED_ENQUEUE_TIME)
+
+    def test_eventhub_trigger_multiple_events_json(self):
+        result = azf_eh.EventHubTriggerConverter.decode(
+            data=self._generate_multiple_iothub_data('json'),
+            trigger_metadata=self._generate_multiple_trigger_metadata()
+        )
+        self.assertIsInstance(result, list)
+        self.assertEqual(len(result), 2)
+
+        self.assertEqual(result[0].enqueued_time, self.MOCKED_ENQUEUE_TIME)
+        self.assertEqual(
+            result[0].get_body().decode('utf-8'), '{"device-status": "good1"}'
+        )
+
+        self.assertEqual(result[1].enqueued_time, self.MOCKED_ENQUEUE_TIME)
+        self.assertEqual(
+            result[1].get_body().decode('utf-8'), '{"device-status": "good2"}'
+        )
+
+    def test_eventhub_trigger_multiple_events_collection_string(self):
+        result = azf_eh.EventHubTriggerConverter.decode(
+            data=self._generate_multiple_iothub_data('collection_string'),
+            trigger_metadata=self._generate_multiple_trigger_metadata()
+        )
+        self.assertIsInstance(result, list)
+        self.assertEqual(len(result), 2)
+
+        self.assertEqual(result[0].enqueued_time, self.MOCKED_ENQUEUE_TIME)
+        self.assertEqual(
+            result[0].get_body().decode('utf-8'), '{"device-status": "good1"}'
+        )
+
+        self.assertEqual(result[1].enqueued_time, self.MOCKED_ENQUEUE_TIME)
+        self.assertEqual(
+            result[1].get_body().decode('utf-8'), '{"device-status": "good2"}'
+        )
+
+    def test_eventhub_trigger_multiple_events_collection_bytes(self):
+        result = azf_eh.EventHubTriggerConverter.decode(
+            data=self._generate_multiple_iothub_data('collection_bytes'),
+            trigger_metadata=self._generate_multiple_trigger_metadata()
+        )
+        self.assertIsInstance(result, list)
+        self.assertEqual(len(result), 2)
+
+        self.assertEqual(result[0].enqueued_time, self.MOCKED_ENQUEUE_TIME)
+        self.assertEqual(
+            result[0].get_body().decode('utf-8'), '{"device-status": "good1"}'
+        )
+
+        self.assertEqual(result[1].enqueued_time, self.MOCKED_ENQUEUE_TIME)
+        self.assertEqual(
+            result[1].get_body().decode('utf-8'), '{"device-status": "good2"}'
+        )
+
+    def test_iothub_metadata_events(self):
+        result = azf_eh.EventHubTriggerConverter.decode(
+            data=self._generate_multiple_iothub_data('json'),
+            trigger_metadata=self._generate_multiple_trigger_metadata()
+        )
+        self.assertIsNotNone(result[0].iothub_metadata)
+        self.assertEqual(
+            result[0].iothub_metadata['connection-device-id'], 'MyTestDevice1'
+        )
+
+        self.assertIsNotNone(result[1].iothub_metadata)
+        self.assertEqual(
+            result[1].iothub_metadata['connection-device-id'], 'MyTestDevice2'
+        )
+
+    def _generate_single_iothub_datum(self, datum_type='json'):
+        datum = '{"device-status": "good"}'
+        if datum_type == 'bytes':
+            datum = datum.encode('utf-8')
+
+        return meta.Datum(datum, datum_type)
+
+    def _generate_multiple_iothub_data(self, data_type='json'):
+        data = '[{"device-status": "good1"}, {"device-status": "good2"}]'
+        if data_type == 'collection_bytes':
+            data = list(
+                map(lambda x: json.dumps(x).encode('utf-8'), json.loads(data))
+            )
+            data = CollectionBytes(data)
+        elif data_type == 'collection_string':
+            data = list(
+                map(lambda x: json.dumps(x), json.loads(data))
+            )
+            data = CollectionString(data)
+
+        return meta.Datum(data, data_type)
+
+    def _generate_single_trigger_metadatum(self):
+        return {
+            'EnqueuedTime': meta.Datum(
+                f'"{self.MOCKED_ENQUEUE_TIME.isoformat()}"', 'json'
+            ),
+            'SystemProperties': meta.Datum(
+                '{"iothub-connection-device-id": "MyTestDevice"}', 'json'
+            )
+        }
+
+    def _generate_multiple_trigger_metadata(self):
+        system_props_array = [
+            {
+                'EnqueuedTimeUtc': self.MOCKED_ENQUEUE_TIME.isoformat(),
+                'iothub-connection-device-id': 'MyTestDevice1',
+            },
+            {
+                'EnqueuedTimeUtc': self.MOCKED_ENQUEUE_TIME.isoformat(),
+                'iothub-connection-device-id': 'MyTestDevice2',
+            }
+        ]
+
+        return {
+            'SystemPropertiesArray': meta.Datum(
+                json.dumps(system_props_array), 'json'
+            )
+        }
