From ef41ec693e0461821038bf00f200891277a2ad13 Mon Sep 17 00:00:00 2001
From: Michael Lavers <kolanos@gmail.com>
Date: Wed, 24 Jul 2019 18:32:05 -0700
Subject: [PATCH] Add post response hook

Closes #338
---
 iopipe/agent.py                    | 13 ++++++++-----
 iopipe/contrib/eventinfo/plugin.py |  3 +++
 iopipe/contrib/logger/plugin.py    |  3 +++
 iopipe/contrib/profiler/plugin.py  |  3 +++
 iopipe/contrib/trace/plugin.py     |  3 +++
 iopipe/plugins.py                  |  4 ++++
 tests/test_plugins.py              |  3 +++
 7 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/iopipe/agent.py b/iopipe/agent.py
index 0fe0117c..d0be53e5 100644
--- a/iopipe/agent.py
+++ b/iopipe/agent.py
@@ -146,11 +146,11 @@ def wrapped(event, context):
             if timeout_duration > 0:
                 logger.debug("Setting timeout duration to %s" % timeout_duration)
 
-            result = None
+            response = None
 
             try:
                 with Timeout(timeout_duration):
-                    result = func(event, context)
+                    response = func(event, context)
             except Exception as e:
                 self.run_hooks("post:invoke", event=event, context=context)
 
@@ -169,6 +169,7 @@ def wrapped(event, context):
                 # event that a timeout occurs and an exception is subsequently raised
                 # within the handler
                 if self.report.sent is False:
+                    self.run_hooks("post:response", response=response)
                     self.report.prepare(e, frame)
                     self.run_hooks("pre:report")
                     self.report.send()
@@ -180,14 +181,15 @@ def wrapped(event, context):
 
                 if context.iopipe.disabled:
                     logger.debug("Reporting disabled for this invocation")
-                    return result
+                    return response
 
+                self.run_hooks("post:response", response=response)
                 self.report.prepare()
                 self.run_hooks("pre:report")
                 self.report.send()
                 self.run_hooks("post:report")
 
-                return result
+                return response
             finally:
                 self.wait_for_futures()
 
@@ -219,7 +221,7 @@ def instantiate(plugin):
 
         return loaded_plugins
 
-    def run_hooks(self, name, event=None, context=None):
+    def run_hooks(self, name, event=None, context=None, response=None):
         """
         Runs plugin hooks for each registered plugin.
         """
@@ -228,6 +230,7 @@ def run_hooks(self, name, event=None, context=None):
             "post:setup": lambda p: p.post_setup(self),
             "pre:invoke": lambda p: p.pre_invoke(event, context),
             "post:invoke": lambda p: p.post_invoke(event, context),
+            "post:response": lambda p: p.post_response(response),
             "pre:report": lambda p: p.pre_report(self.report),
             "post:report": lambda p: p.post_report(self.report),
         }
diff --git a/iopipe/contrib/eventinfo/plugin.py b/iopipe/contrib/eventinfo/plugin.py
index a68ff83a..e5d0ae9f 100644
--- a/iopipe/contrib/eventinfo/plugin.py
+++ b/iopipe/contrib/eventinfo/plugin.py
@@ -41,6 +41,9 @@ def post_invoke(self, event, context):
         if self.enabled:
             metrics_for_event_type(event, context)
 
+    def post_response(self, response):
+        pass
+
     def pre_report(self, report):
         pass
 
diff --git a/iopipe/contrib/logger/plugin.py b/iopipe/contrib/logger/plugin.py
index 750deac6..932144f8 100644
--- a/iopipe/contrib/logger/plugin.py
+++ b/iopipe/contrib/logger/plugin.py
@@ -110,6 +110,9 @@ def post_invoke(self, event, context):
         if self.enabled and self.redirect_stdout is True:
             sys.stdout = sys.__stdout__
 
+    def post_response(self, response):
+        pass
+
     def pre_report(self, report):
         if self.enabled:
             self.handler.flush()
diff --git a/iopipe/contrib/profiler/plugin.py b/iopipe/contrib/profiler/plugin.py
index 4989800a..e3258d0e 100644
--- a/iopipe/contrib/profiler/plugin.py
+++ b/iopipe/contrib/profiler/plugin.py
@@ -63,6 +63,9 @@ def post_invoke(self, event, context):
             self.profile.disable()
             self.context.iopipe.label("@iopipe/plugin-profiler")
 
+    def post_response(self, response):
+        pass
+
     def pre_report(self, report):
         if self.profile is not None:
             if self.signed_request is not None:
diff --git a/iopipe/contrib/trace/plugin.py b/iopipe/contrib/trace/plugin.py
index b42173a8..36e30b63 100644
--- a/iopipe/contrib/trace/plugin.py
+++ b/iopipe/contrib/trace/plugin.py
@@ -58,6 +58,9 @@ def post_invoke(self, event, context):
         if self.auto_http is True:
             restore_http_requests()
 
+    def post_response(self, response):
+        pass
+
     def pre_report(self, report):
         if self.auto_measure:
             add_timeline_measures(self.timeline)
diff --git a/iopipe/plugins.py b/iopipe/plugins.py
index 2efcde2d..8d464216 100644
--- a/iopipe/plugins.py
+++ b/iopipe/plugins.py
@@ -79,6 +79,10 @@ def pre_invoke(self, event, context):
     def post_invoke(self, event, context):
         return NotImplemented
 
+    @abc.abstractmethod
+    def post_response(self, response):
+        return NotImplemented
+
     @abc.abstractmethod
     def pre_report(self, report):
         return NotImplemented
diff --git a/tests/test_plugins.py b/tests/test_plugins.py
index 2411f5c5..1c17f20a 100644
--- a/tests/test_plugins.py
+++ b/tests/test_plugins.py
@@ -50,6 +50,9 @@ def pre_invoke(self, event, context):
         def post_invoke(self, event, context):
             pass
 
+        def post_response(self, response):
+            pass
+
         def pre_report(self, report):
             pass