Allow disabling of single invocation reporting (#324)

* Allow disabling of single invocation reporting

Closes #289

* Add usage to README

README.md

@@ -22,6 +22,7 @@ This package provides analytics and distributed tracing for event-driven applica
   - [Custom Metrics](https://github.com/iopipe/iopipe-python#custom-metrics)
   - [Labels](https://github.com/iopipe/iopipe-python#labels)
   - [Core Agent](https://github.com/iopipe/iopipe-python#core-agent)
+  - [Disabling Reporting](https://github.com/iopipe/iopipe-python#disabling-reporting)
 - [Plugins](https://github.com/iopipe/iopipe-python#plugins)
   - [Event Info Plugin](https://github.com/iopipe/iopipe-python#event-info-plugin)
   - [Logger Plugin](https://github.com/iopipe/iopipe-python#logger-plugin)
@@ -79,7 +80,7 @@ iopipe = IOpipe('your project token here')
 
 @iopipe
 def handler(event, context):
-  pass
+    pass
 ```
 
  The agent comes preloaded with the [Event Info](https://github.com/iopipe/iopipe-python#event-info-plugin), [Profiler](https://github.com/iopipe/iopipe-python#profiler-plugin) and [Trace](https://github.com/iopipe/iopipe-python#trace-plugin) plugins. See the relevant plugin sections for usage.
@@ -121,17 +122,17 @@ iopipe = IOpipe()
 
 @iopipe
 def handler(event, context):
-  raise Exception('This exception will be added to the IOpipe report automatically')
+    raise Exception('This exception will be added to the IOpipe report automatically')
 
 # Example 2: caught exceptions
 
 @iopipe
 def handler(event, context):
-  try:
-    raise Exception('This exception is being caught by your function')
-  except Exception as e:
-    # Makes sure the exception is added to the report
-    context.iopipe.error(e)
+    try:
+        raise Exception('This exception is being caught by your function')
+    except Exception as e:
+        # Makes sure the exception is added to the report
+        context.iopipe.error(e)
 ```
 
 It is important to note that a report is sent to IOpipe when `error()` is called. So you should only record exceptions this way for failure states. For caught exceptions that are not a failure state, it is recommended to use custom metrics (see below).
@@ -147,9 +148,9 @@ iopipe = IOpipe()
 
 @iopipe
 def handler(event, context):
-  # the name of the metric must be a string
-  # numerical (int, long, float) and string types supported for values
-  context.iopipe.metric('my_metric', 42)
+    # the name of the metric must be a string
+    # numerical (int, long, float) and string types supported for values
+    context.iopipe.metric('my_metric', 42)
 ```
 
 Metric key names are limited to 128 characters, and string values are limited to 1024 characters.
@@ -165,8 +166,8 @@ iopipe = IOpipe()
 
 @iopipe
 def handler(event, context):
-  # the name of the tag must be a string
-  context.iopipe.label('this-invocation-is-special')
+    # the name of the label must be a string
+    context.iopipe.label('this-invocation-is-special')
 ```
 
 ### Core Agent
@@ -182,9 +183,26 @@ iopipe = IOpipeCore(plugins=[TracePlugin()])
 
 @iopipe
 def handler(event, context):
-pass
+    pass
 ```
 
+### Disabling Reporting
+
+You can programmatically disable IOpipe reporting for a single invocation using the `disable` method:
+
+```python
+from iopipe import IOpipe
+
+iopipe = IOpipe()
+
+@iopipe
+def handler(event, context):
+    if some_condition:
+        context.iopipe.disable()
+```
+
+Reporting will be re-enabled on the next invocation.
+
 ## Plugins
 
 IOpipe's functionality can be extended through plugins. Plugins hook into the agent lifecycle to allow you to perform additional analytics.

iopipe/agent.py

@@ -146,6 +146,10 @@ def wrapped(event, context):
             except Exception as e:
                 self.run_hooks("post:invoke", event=event, context=context)
 
+                if context.iopipe.disabled:
+                    logger.debug("Reporting disabled for this invocation")
+                    raise e
+
                 frame = None
                 if isinstance(e, TimeoutError):
                     logger.debug(
@@ -165,15 +169,20 @@ def wrapped(event, context):
                 raise e
             else:
                 self.run_hooks("post:invoke", event=event, context=context)
+
+                if context.iopipe.disabled:
+                    logger.debug("Reporting disabled for this invocation")
+                    return result
+
                 self.report.prepare()
                 self.run_hooks("pre:report")
                 self.report.send()
                 self.run_hooks("post:report")
+
+                return result
             finally:
                 self.wait_for_futures()
 
-            return result
-
         return wrapped
 
     decorator = __call__

iopipe/context.py

@@ -38,6 +38,7 @@ class IOpipeContext(object):
     def __init__(self, instance):
         self.instance = instance
         self.log = LogWrapper(self)
+        self.disabled = False
 
     def metric(self, key, value):
         if self.instance.report is None:
@@ -125,3 +126,6 @@ def register(self, name, value, force=False):
     def unregister(self, name):
         if hasattr(self, name):
             delattr(self, name)
+
+    def disable(self):
+        self.disabled = True

tests/conftest.py

@@ -125,7 +125,7 @@ def mock_context():
 
 @pytest.fixture
 def handler(iopipe):
-    @iopipe.decorator
+    @iopipe
     def _handler(event, context):
         pass
 
@@ -134,7 +134,7 @@ def _handler(event, context):
 
 @pytest.fixture
 def handler_with_events(iopipe):
-    @iopipe.decorator
+    @iopipe
     def _handler_with_events(event, context):
         iopipe.log("key1", 2)
         iopipe.log("key2", "qualitative value")
@@ -150,7 +150,7 @@ def _handler_with_events(event, context):
 
 @pytest.fixture
 def handler_with_labels(iopipe):
-    @iopipe.decorator
+    @iopipe
     def _handler_with_labels(event, context):
         context.iopipe.label("a-label")
         context.iopipe.label("a-label")  # duplicates are dropped
@@ -165,7 +165,7 @@ def _handler_with_labels(event, context):
 
 @pytest.fixture
 def handler_that_errors(iopipe):
-    @iopipe.decorator
+    @iopipe
     def _handler_that_errors(event, context):
         raise ValueError("Behold, a value error")
 
@@ -174,7 +174,7 @@ def _handler_that_errors(event, context):
 
 @pytest.fixture
 def handler_that_timeouts(iopipe):
-    @iopipe.decorator
+    @iopipe
     def _handler_that_timeouts(event, context):
         time.sleep(1)
         raise Exception("Should timeout before this is raised")
@@ -189,3 +189,22 @@ def _handler(event, context):
         pass
 
     return iopipe_with_sync_http, _handler
+
+
+@pytest.fixture
+def handler_that_disables_reporting(iopipe):
+    @iopipe
+    def _handler_that_disables_reporting(event, context):
+        context.iopipe.disable()
+
+    return iopipe, _handler_that_disables_reporting
+
+
+@pytest.fixture
+def handler_that_disables_reporting_with_error(iopipe):
+    @iopipe
+    def _handler_that_disables_reporting_with_error(event, context):
+        context.iopipe.disable()
+        raise Exception("An error happened")
+
+    return iopipe, _handler_that_disables_reporting_with_error

tests/test_agent.py

@@ -1,4 +1,5 @@
 import mock
+import pytest
 
 from iopipe import constants
 
@@ -65,10 +66,8 @@ def test_erroring(mock_send_report, handler_that_errors, mock_context):
     """Assert that the agent catches and traces uncaught exceptions"""
     iopipe, handler = handler_that_errors
 
-    try:
+    with pytest.raises(ValueError):
         handler(None, mock_context)
-    except Exception:
-        pass
 
     assert iopipe.report.report["errors"]["name"] == "ValueError"
     assert iopipe.report.report["errors"]["message"] == "Behold, a value error"
@@ -101,10 +100,8 @@ def test_timeouts_disable(mock_send_report, handler_that_timeouts, mock_context)
     # The default is 0.15, so 150 / 1000 - 0.15 = 0
     mock_context.set_remaining_time_in_millis(150)
 
-    try:
+    with pytest.raises(Exception):
         handler(None, mock_context)
-    except Exception:
-        pass
 
     # Exception will occur because timeout is disabled
     assert iopipe.report.report["errors"] != {}
@@ -135,3 +132,23 @@ class InvalidContext(object):
     invalid_context = InvalidContext()
 
     assert iopipe.validate_context(invalid_context) is False
+
+
+def test_disabled_reporting(handler_that_disables_reporting, mock_context):
+    """Assert that reporting is disabled for an invocation"""
+    iopipe, handler = handler_that_disables_reporting
+    handler(None, mock_context)
+
+    assert iopipe.report.sent is False
+
+
+def test_disabled_reporting_with_error(
+    handler_that_disables_reporting_with_error, mock_context
+):
+    """Assert that reporting is disabled for an invocation with error"""
+    iopipe, handler = handler_that_disables_reporting_with_error
+
+    with pytest.raises(Exception):
+        handler(None, mock_context)
+
+    assert iopipe.report.sent is False