feat(greengrasssdk.py): add greengrass support (#289)

Author: ranrib

README.md

@@ -443,6 +443,7 @@ Epsagon provides out-of-the-box instrumentation (tracing) for many popular frame
 |azure.cosmos        |`>=4.0.0`                  |
 |celery              |`>=4.0.0`                  |
 |grpc                |`>=0.3-10`                 |
+|greengrasssdk       |`>=1.4.0`                 |
 |SQLAlchemy                |`>=1.2.0`                 |
 
 

epsagon/events/greengrasssdk.py

@@ -0,0 +1,77 @@
+"""
+Greengrass events module.
+"""
+
+from __future__ import absolute_import
+import traceback
+from uuid import uuid4
+
+from epsagon.utils import add_data_if_needed
+from ..trace import trace_factory
+from ..event import BaseEvent
+
+
+class GreengrassPublishEvent(BaseEvent):
+    """
+    Represents Greengrass publish event.
+    """
+
+    ORIGIN = 'greengrasssdk'
+    RESOURCE_TYPE = 'greengrass'
+
+    # pylint: disable=W0613
+    def __init__(self, wrapped, instance, args, kwargs, start_time, response,
+                 exception):
+        """
+        Initialize.
+        :param wrapped: wrapt's wrapped
+        :param instance: wrapt's instance
+        :param args: wrapt's args
+        :param kwargs: wrapt's kwargs
+        :param start_time: Start timestamp (epoch)
+        :param response: response data
+        :param exception: Exception (if happened)
+        """
+        super(GreengrassPublishEvent, self).__init__(start_time)
+
+        self.event_id = 'greengrass-{}'.format(str(uuid4()))
+
+        self.resource['name'] = kwargs.get('topic', 'N/A')
+        self.resource['operation'] = 'publish'
+        if kwargs.get('queueFullPolicy'):
+            self.resource['metadata']['aws.greengrass.queueFullPolicy'] = (
+                kwargs.get('queueFullPolicy')
+            )
+
+        add_data_if_needed(
+            self.resource['metadata'],
+            'aws.greengrass.payload',
+            kwargs.get('payload')
+        )
+
+        if exception is not None:
+            self.set_exception(exception, traceback.format_exc())
+
+
+class GreengrassEventFactory(object):
+    """
+    Factory class, generates Greengrass event.
+    """
+
+    @staticmethod
+    def create_event(wrapped, instance, args, kwargs, start_time, response,
+                     exception):
+        """
+        Create an event according to the given api_name.
+        """
+        event = GreengrassPublishEvent(
+            wrapped,
+            instance,
+            args,
+            kwargs,
+            start_time,
+            response,
+            exception
+        )
+
+        trace_factory.add_event(event)

epsagon/modules/greengrasssdk.py

@@ -0,0 +1,32 @@
+"""
+greengrasssdk patcher module.
+"""
+
+from __future__ import absolute_import
+import wrapt
+from epsagon.modules.general_wrapper import wrapper
+from ..events.greengrasssdk import GreengrassEventFactory
+
+
+def _wrapper(wrapped, instance, args, kwargs):
+    """
+    General wrapper for greengrasssdk instrumentation.
+    :param wrapped: wrapt's wrapped
+    :param instance: wrapt's instance
+    :param args: wrapt's args
+    :param kwargs: wrapt's kwargs
+    :return: None
+    """
+    return wrapper(GreengrassEventFactory, wrapped, instance, args, kwargs)
+
+
+def patch():
+    """
+    Patch module.
+    :return: None
+    """
+    wrapt.wrap_function_wrapper(
+        'greengrasssdk.IoTDataPlane',
+        'Client.publish',
+        _wrapper
+    )

tests/modules/test_greengrasssdk.py

@@ -0,0 +1,19 @@
+import mock
+from epsagon.events.greengrasssdk import GreengrassEventFactory
+
+
+@mock.patch('epsagon.trace.TraceFactory.add_event')
+def test_sanity(add_event_mock):
+    params = {
+        'topic': 'name',
+        'queueFullPolicy': True,
+        'payload': 'test',
+    }
+    GreengrassEventFactory.create_event(None, None, None, params, None, None, None)
+    add_event_mock.assert_called_once()
+    event = add_event_mock.call_args_list[0].args[0]
+    assert event.event_id.startswith('greengrass-')
+    assert event.resource['name'] == 'name'
+    assert event.resource['operation'] == 'publish'
+    assert event.resource['metadata']['aws.greengrass.queueFullPolicy'] == True
+    assert event.resource['metadata']['aws.greengrass.payload'] == 'test'