From 6e0c738711ab73d1413b732fa6984a028694ea5a Mon Sep 17 00:00:00 2001
From: Michael Lavers <kolanos@gmail.com>
Date: Thu, 25 Jul 2019 12:39:33 -0700
Subject: [PATCH 1/5] Response Event Info, Step Function Meta Data

Closes #340
---
 .circleci/config.yml                    |  6 +++---
 .pre-commit-config.yaml                 |  9 +++------
 README.md                               | 16 ++++++++++------
 iopipe/agent.py                         |  7 +++++--
 iopipe/context.py                       | 18 +++++++++++++++---
 iopipe/contrib/eventinfo/event_types.py |  7 +++++++
 iopipe/contrib/eventinfo/plugin.py      |  5 +++--
 pyproject.toml                          | 18 ++++++++++++++++++
 setup.py                                |  2 +-
 tests/contrib/logger/conftest.py        |  2 +-
 tests/contrib/logger/test_plugin.py     |  2 +-
 11 files changed, 67 insertions(+), 25 deletions(-)
 create mode 100644 pyproject.toml

diff --git a/.circleci/config.yml b/.circleci/config.yml
index ce208bde..1f5d2221 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -47,9 +47,9 @@ jobs:
       - run:
           name: Check code style
           command: |
-            pip install black==18.6b2
-            black --check --line-length=88 --safe iopipe
-            black --check --line-length=88 --safe tests
+            pip install black==19.3b0
+            black iopipe
+            black tests
 
   coverage:
     working_directory: ~/iopipe-python
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 6f509e6a..3bc19511 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -1,8 +1,5 @@
-repos:
--   repo: https://github.com/ambv/black
-    rev: 18.6b2
-
+-   repo: https://github.com/psf/black
+    rev: stable
     hooks:
     - id: black
-      args: [--line-length=88, --safe]
-      python_version: python3.6
+      language_version: python3.7
diff --git a/README.md b/README.md
index 1953bebf..aa1060dd 100644
--- a/README.md
+++ b/README.md
@@ -497,6 +497,9 @@ class MyPlugin(Plugin):
     def post_invoke(self, event, context):
         pass
 
+    def post_response(self, response):
+        pass
+
     def pre_report(self, report):
         pass
 
@@ -519,12 +522,13 @@ A plugin has the following properties defined:
 
 A plugin has the following methods defined:
 
-- `pre_setup`: Is called once prior to the agent initialization. Is passed the `iopipe` instance.
-- `post_setup`: Is called once after the agent is initialized, is passed the `iopipe` instance.
-- `pre_invoke`: Is called prior to each invocation, is passed the `event` and `context` of the invocation.
-- `post_invoke`: Is called after each invocation, is passed the `event` and `context` of the invocation.
-- `pre_report`: Is called prior to each report being sent, is passed the `report` instance.
-- `post_report`: Is called after each report is sent, is passed the `report` instance.
+- `pre_setup`: Is called once prior to the agent initialization; is passed the `iopipe` instance.
+- `post_setup`: Is called once after the agent is initialized; is passed the `iopipe` instance.
+- `pre_invoke`: Is called prior to each invocation; is passed the `event` and `context` of the invocation.
+- `post_invoke`: Is called after each invocation; is passed the `event` and `context` of the invocation.
+- `post_response`: Is called after the invocation response; is passed the `response`value.
+- `pre_report`: Is called prior to each report being sent; is passed the `report` instance.
+- `post_report`: Is called after each report is sent; is passed the `report` instance.
 
 ## Supported Python Versions
 
diff --git a/iopipe/agent.py b/iopipe/agent.py
index d0be53e5..68cc015c 100644
--- a/iopipe/agent.py
+++ b/iopipe/agent.py
@@ -84,7 +84,7 @@ def error(self, error):
 
     err = error
 
-    def __call__(self, func):
+    def __call__(self, func, **kwargs):
         @functools.wraps(func)
         def wrapped(event, context):
             # Skip if function is already instrumented
@@ -96,7 +96,7 @@ def wrapped(event, context):
 
             logger.debug("Wrapping %s with IOpipe decorator" % repr(func))
 
-            self.context = context = ContextWrapper(context, self)
+            self.context = context = ContextWrapper(context, self, **kwargs)
 
             # if env var IOPIPE_ENABLED is set to False skip reporting
             if self.config["enabled"] is False:
@@ -197,6 +197,9 @@ def wrapped(event, context):
 
     decorator = __call__
 
+    def step(self, func):
+        return self(func, step_function=True)
+
     def load_plugins(self, plugins):
         """
         Loads plugins that match the `Plugin` interface and are instantiated.
diff --git a/iopipe/context.py b/iopipe/context.py
index de41245a..d34bc654 100644
--- a/iopipe/context.py
+++ b/iopipe/context.py
@@ -1,6 +1,7 @@
 import decimal
 import logging
 import numbers
+import uuid
 import warnings
 
 from . import constants
@@ -25,20 +26,22 @@ def __getattr__(self, name):
 
 
 class ContextWrapper(object):
-    def __init__(self, base_context, instance):
+    def __init__(self, base_context, instance, **kwargs):
         self.base_context = base_context
         self.instance = instance
-        self.iopipe = IOpipeContext(self.instance)
+        self.iopipe = IOpipeContext(self.instance, **kwargs)
 
     def __getattr__(self, name):
         return getattr(self.base_context, name)
 
 
 class IOpipeContext(object):
-    def __init__(self, instance):
+    def __init__(self, instance, **kwargs):
         self.instance = instance
         self.log = LogWrapper(self)
         self.disabled = False
+        self.is_step_function = kwargs.pop("step_function", False)
+        self.step_meta = None
 
     def metric(self, key, value):
         if self.instance.report is None:
@@ -129,3 +132,12 @@ def unregister(self, name):
 
     def disable(self):
         self.disabled = True
+
+    def collect_step_meta(self, event):
+        if self.is_step_function:
+            self.step_meta = event.get("iopipe", {"id": str(uuid.uuid4()), "step": 0})
+
+    def inject_step_meta(self, response):
+        if self.step_meta and isinstance(response, dict):
+            self.step_meta["step"] += 1
+            response["iopipe"] = self.step_meta
diff --git a/iopipe/contrib/eventinfo/event_types.py b/iopipe/contrib/eventinfo/event_types.py
index ce1ba957..19ba2c65 100644
--- a/iopipe/contrib/eventinfo/event_types.py
+++ b/iopipe/contrib/eventinfo/event_types.py
@@ -3,6 +3,7 @@
 
 class EventType(object):
     keys = []
+    response_keys = []
     exclude_keys = []
     required_keys = []
     source = None
@@ -266,3 +267,9 @@ def metrics_for_event_type(event, context):
             event_info = event_type.collect()
             [context.iopipe.metric(k, v) for k, v in event_info.items()]
             break
+
+    if context.iopipe.is_step_function:
+        context.iopipe.collect_step_meta(event)
+        if context.iopipe.step_meta:
+            for key, value in context.iopipe.step_meta.items():
+                context.iopipe.metric("@iopipe/event-info.stepFunction.%s" % key, value)
diff --git a/iopipe/contrib/eventinfo/plugin.py b/iopipe/contrib/eventinfo/plugin.py
index e5d0ae9f..33ae252e 100644
--- a/iopipe/contrib/eventinfo/plugin.py
+++ b/iopipe/contrib/eventinfo/plugin.py
@@ -35,14 +35,15 @@ def post_setup(self, iopipe):
         pass
 
     def pre_invoke(self, event, context):
-        pass
+        self.context = context
 
     def post_invoke(self, event, context):
         if self.enabled:
             metrics_for_event_type(event, context)
 
     def post_response(self, response):
-        pass
+        if self.enabled:
+            self.context.iopipe.inject_step_meta(response)
 
     def pre_report(self, report):
         pass
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 00000000..d05520fe
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,18 @@
+[tool.black]
+line-length = 88
+target-version = ['py27', 'py36', 'py37', 'py38']
+include = '\.pyi?$'
+exclude = '''
+/(
+    \.eggs
+  | \.git
+  | \.hg
+  | \.mypy_cache
+  | \.tox
+  | \.venv
+  | _build
+  | buck-out
+  | build
+  | dist
+)/
+'''
diff --git a/setup.py b/setup.py
index 4fbc6114..1c7061d2 100644
--- a/setup.py
+++ b/setup.py
@@ -39,7 +39,7 @@
     packages=find_packages(exclude=("tests", "tests.*")),
     extras_require={
         "coverage": coverage_requires,
-        "dev": tests_require + ["black==18.6b2", "pre-commit"],
+        "dev": tests_require + ["black==19.3b0", "pre-commit"],
     },
     install_requires=install_requires,
     setup_requires=["pytest-runner==4.2"],
diff --git a/tests/contrib/logger/conftest.py b/tests/contrib/logger/conftest.py
index e7f0be72..82ff1a1d 100644
--- a/tests/contrib/logger/conftest.py
+++ b/tests/contrib/logger/conftest.py
@@ -29,7 +29,7 @@ def _handler(event, context):
         except Exception as e:
             context.iopipe.log.exception(e)
 
-        print("This is not a misprint.")
+        print ("This is not a misprint.")
 
     return iopipe_with_logger, _handler
 
diff --git a/tests/contrib/logger/test_plugin.py b/tests/contrib/logger/test_plugin.py
index 9fbae015..fa20ffaa 100644
--- a/tests/contrib/logger/test_plugin.py
+++ b/tests/contrib/logger/test_plugin.py
@@ -161,7 +161,7 @@ def test__logger_plugin__use_tmp__disk_used(
         pytest.skip("this test requires linux, skipping")
 
     disk_usage = read_disk()
-    print(disk_usage)
+    print (disk_usage)
 
     iopipe, handler = handler_with_logger_use_tmp
     handler({}, mock_context)

From 0907562a0dc15b6438a6ce8eecc01e463b70bd27 Mon Sep 17 00:00:00 2001
From: Michael Lavers <kolanos@gmail.com>
Date: Thu, 25 Jul 2019 13:28:35 -0700
Subject: [PATCH 2/5] Adding tests

---
 iopipe/context.py                      |  5 +++--
 tests/conftest.py                      | 13 +++++++++++++
 tests/contrib/eventinfo/conftest.py    | 10 ++++++++++
 tests/contrib/eventinfo/test_plugin.py | 23 +++++++++++++++++++++++
 tests/test_agent.py                    | 10 ++++++++++
 5 files changed, 59 insertions(+), 2 deletions(-)

diff --git a/iopipe/context.py b/iopipe/context.py
index d34bc654..bdfd33f1 100644
--- a/iopipe/context.py
+++ b/iopipe/context.py
@@ -139,5 +139,6 @@ def collect_step_meta(self, event):
 
     def inject_step_meta(self, response):
         if self.step_meta and isinstance(response, dict):
-            self.step_meta["step"] += 1
-            response["iopipe"] = self.step_meta
+            step_meta = self.step_meta.copy()
+            step_meta["step"] += 1
+            response["iopipe"] = step_meta
diff --git a/tests/conftest.py b/tests/conftest.py
index c30ffaf8..81c2757c 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -208,3 +208,16 @@ def _handler_that_disables_reporting_with_error(event, context):
         raise Exception("An error happened")
 
     return iopipe, _handler_that_disables_reporting_with_error
+
+
+@pytest.fixture
+def handler_step_function(iopipe):
+    @iopipe
+    def _handler_not_step_function(event, context):
+        assert context.iopipe.is_step_function is False
+
+    @iopipe.step
+    def _handler_step_function(event, context):
+        assert context.iopipe.is_step_function is True
+
+    return iopipe, _handler_step_function
diff --git a/tests/contrib/eventinfo/conftest.py b/tests/contrib/eventinfo/conftest.py
index 83c01eff..1c2c6280 100644
--- a/tests/contrib/eventinfo/conftest.py
+++ b/tests/contrib/eventinfo/conftest.py
@@ -26,6 +26,16 @@ def _handler(event, context):
     return iopipe_with_eventinfo, _handler
 
 
+@pytest.fixture
+def handler_step_function_with_eventinfo(iopipe_with_eventinfo):
+    @iopipe_with_eventinfo.step
+    def _handler(event, context):
+        assert context.iopipe.is_step_function is True
+        return {}
+
+    return iopipe_with_eventinfo, _handler
+
+
 def _load_event(name):
     json_file = os.path.join(os.path.dirname(__file__), "events", "%s.json" % name)
     with open(json_file) as f:
diff --git a/tests/contrib/eventinfo/test_plugin.py b/tests/contrib/eventinfo/test_plugin.py
index 5af58a51..d3022a89 100644
--- a/tests/contrib/eventinfo/test_plugin.py
+++ b/tests/contrib/eventinfo/test_plugin.py
@@ -112,3 +112,26 @@ def test__eventinfo_plugin__enabled(monkeypatch):
     plugin = EventInfoPlugin(enabled=False)
 
     assert plugin.enabled is True
+
+
+@mock.patch("iopipe.report.send_report", autospec=True)
+def test__eventinfo_plugin__step_function(
+    mock_send_report, handler_step_function_with_eventinfo, event_apigw, mock_context
+):
+    iopipe, handler = handler_step_function_with_eventinfo
+
+    plugins = iopipe.config["plugins"]
+    assert len(plugins) == 1
+    assert plugins[0].enabled is True
+    assert plugins[0].name == "event-info"
+
+    response1 = handler(event_apigw, mock_context)
+    assert "iopipe" in response1
+    assert "id" in response1["iopipe"]
+    assert "step" in response1["iopipe"]
+
+    response2 = handler(response1, mock_context)
+
+    assert "iopipe" in response2
+    assert response1["iopipe"]["id"] == response2["iopipe"]["id"]
+    assert response2["iopipe"]["step"] > response1["iopipe"]["step"]
diff --git a/tests/test_agent.py b/tests/test_agent.py
index e798ddce..317c19f8 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -177,3 +177,13 @@ def test_double_instrumentation(mock_send_report, handler, mock_context, monkeyp
     double_wrapped(None, mock_context)
 
     assert mock_send_report.call_count == 1
+
+
+@mock.patch("iopipe.report.send_report", autospec=True)
+def test_step_function(mock_send_report, handler_step_function, mock_context):
+    """Assert that step functions are identified as such"""
+    iopipe, handler = handler_step_function
+
+    handler(None, mock_context)
+
+    assert iopipe.context.iopipe.is_step_function is True

From c55f21fc9055a028de7f8a71bb7f8b0b9914778d Mon Sep 17 00:00:00 2001
From: Michael Lavers <kolanos@gmail.com>
Date: Thu, 25 Jul 2019 14:18:03 -0700
Subject: [PATCH 3/5] Add LambdaProxy response type

---
 iopipe/contrib/eventinfo/event_types.py    | 14 +++++++++++
 iopipe/contrib/eventinfo/plugin.py         |  5 ++--
 iopipe/contrib/eventinfo/response_types.py | 28 ++++++++++++++++++++++
 3 files changed, 45 insertions(+), 2 deletions(-)
 create mode 100644 iopipe/contrib/eventinfo/response_types.py

diff --git a/iopipe/contrib/eventinfo/event_types.py b/iopipe/contrib/eventinfo/event_types.py
index 19ba2c65..618c1611 100644
--- a/iopipe/contrib/eventinfo/event_types.py
+++ b/iopipe/contrib/eventinfo/event_types.py
@@ -1,3 +1,4 @@
+from .response_types import LambdaProxy
 from .util import collect_all_keys, get_value, has_key, slugify
 
 
@@ -6,10 +7,13 @@ class EventType(object):
     response_keys = []
     exclude_keys = []
     required_keys = []
+    response_type = None
     source = None
 
     def __init__(self, event):
         self.event = event
+        if self.response_type:
+            self.response_type = self.response_type(self.type)
 
     @property
     def slug(self):
@@ -84,6 +88,7 @@ class ApiGateway(EventType):
         "resource",
     ]
     required_keys = ["headers", "httpMethod", "path", "requestContext", "resource"]
+    response_type = LambdaProxy
 
 
 class CloudFront(EventType):
@@ -173,6 +178,7 @@ class ServerlessLambda(EventType):
         ("stage", "requestContext.stage"),
     ]
     required_keys = ["identity.userAgent", "identity.sourceIp", "identity.accountId"]
+    response_type = LambdaProxy
 
 
 class SES(EventType):
@@ -273,3 +279,11 @@ def metrics_for_event_type(event, context):
         if context.iopipe.step_meta:
             for key, value in context.iopipe.step_meta.items():
                 context.iopipe.metric("@iopipe/event-info.stepFunction.%s" % key, value)
+
+    return event_type
+
+
+def metrics_for_response_type(event_type, context, response):
+    if event_type.response_type:
+        response_info = event_type.response_type.collect(response)
+        [context.iopipe.metric(k, v) for k, v in response_info.items()]
diff --git a/iopipe/contrib/eventinfo/plugin.py b/iopipe/contrib/eventinfo/plugin.py
index 33ae252e..3f4c5182 100644
--- a/iopipe/contrib/eventinfo/plugin.py
+++ b/iopipe/contrib/eventinfo/plugin.py
@@ -3,7 +3,7 @@
 
 from iopipe.plugins import Plugin
 
-from .event_types import metrics_for_event_type
+from .event_types import metrics_for_event_type, metrics_for_response_type
 
 
 class EventInfoPlugin(Plugin):
@@ -39,10 +39,11 @@ def pre_invoke(self, event, context):
 
     def post_invoke(self, event, context):
         if self.enabled:
-            metrics_for_event_type(event, context)
+            self.event_type = metrics_for_event_type(event, context)
 
     def post_response(self, response):
         if self.enabled:
+            metrics_for_response_type(self.event_type, self.context, response)
             self.context.iopipe.inject_step_meta(response)
 
     def pre_report(self, report):
diff --git a/iopipe/contrib/eventinfo/response_types.py b/iopipe/contrib/eventinfo/response_types.py
new file mode 100644
index 00000000..821c4a38
--- /dev/null
+++ b/iopipe/contrib/eventinfo/response_types.py
@@ -0,0 +1,28 @@
+from .util import get_value
+
+
+class ResponseType(object):
+    keys = []
+
+    def __init__(self, event_type):
+        self.event_type = event_type
+
+    def collect(self, response):
+        response_info = {}
+
+        for key in self.keys:
+            if isinstance(key, tuple):
+                old_key, new_key = key
+            else:
+                old_key = new_key = key
+            value = get_value(response, old_key)
+            if value is not None:
+                response_info[
+                    "@iopipe/event-info.%s.response.%s" % (self.event_type, new_key)
+                ] = value
+
+        return response_info
+
+
+class LambdaProxy(ResponseType):
+    keys = ["statusCode"]

From e5ca6952d88028987ff3a32944404fdfe0fb528d Mon Sep 17 00:00:00 2001
From: Michael Lavers <kolanos@gmail.com>
Date: Fri, 26 Jul 2019 09:50:38 -0700
Subject: [PATCH 4/5] Add test for http response collection

---
 tests/contrib/eventinfo/conftest.py    |  9 +++++++++
 tests/contrib/eventinfo/test_plugin.py | 26 ++++++++++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/tests/contrib/eventinfo/conftest.py b/tests/contrib/eventinfo/conftest.py
index 1c2c6280..c5742457 100644
--- a/tests/contrib/eventinfo/conftest.py
+++ b/tests/contrib/eventinfo/conftest.py
@@ -36,6 +36,15 @@ def _handler(event, context):
     return iopipe_with_eventinfo, _handler
 
 
+@pytest.fixture
+def handler_http_response_with_eventinfo(iopipe_with_eventinfo):
+    @iopipe_with_eventinfo
+    def _handler(event, context):
+        return {"statusCode": 200, "body": "success"}
+
+    return iopipe_with_eventinfo, _handler
+
+
 def _load_event(name):
     json_file = os.path.join(os.path.dirname(__file__), "events", "%s.json" % name)
     with open(json_file) as f:
diff --git a/tests/contrib/eventinfo/test_plugin.py b/tests/contrib/eventinfo/test_plugin.py
index d3022a89..f036222b 100644
--- a/tests/contrib/eventinfo/test_plugin.py
+++ b/tests/contrib/eventinfo/test_plugin.py
@@ -135,3 +135,29 @@ def test__eventinfo_plugin__step_function(
     assert "iopipe" in response2
     assert response1["iopipe"]["id"] == response2["iopipe"]["id"]
     assert response2["iopipe"]["step"] > response1["iopipe"]["step"]
+
+
+@mock.patch("iopipe.report.send_report", autospec=True)
+def test__eventinfo_plugin__http_response(
+    mock_send_report, handler_http_response_with_eventinfo, event_apigw, mock_context
+):
+    iopipe, handler = handler_http_response_with_eventinfo
+
+    handler(event_apigw, mock_context)
+    metrics = iopipe.report.custom_metrics
+
+    assert any(
+        (
+            m["name"] == "@iopipe/event-info.apiGateway.response.statusCode"
+            for m in metrics
+        )
+    )
+
+    metric = next(
+        (
+            m
+            for m in metrics
+            if m["name"] == "@iopipe/event-info.apiGateway.response.statusCode"
+        )
+    )
+    assert metric["n"] == 200

From 5377b62c7142782be5956ff569ad0efd685de036 Mon Sep 17 00:00:00 2001
From: Michael Lavers <kolanos@gmail.com>
Date: Fri, 26 Jul 2019 10:02:17 -0700
Subject: [PATCH 5/5] Update README

---
 README.md | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/README.md b/README.md
index aa1060dd..fe7f26c2 100644
--- a/README.md
+++ b/README.md
@@ -23,6 +23,7 @@ This package provides analytics and distributed tracing for event-driven applica
   - [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)
+  - [Step Functions](https://github.com/iopipe/iopipe-python#step-functions)
 - [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)
@@ -205,6 +206,22 @@ def handler(event, context):
 
 Reporting will be re-enabled on the next invocation.
 
+### Step Functions
+
+IOpipe is compatible with AWS Lambda step functions. To enable step function tracing:
+
+```python
+from iopipe import IOpipe
+
+iopipe = IOpipe()
+
+@iopipe.step
+def handler(event, context):
+    pass
+```
+
+The `@iopipe.step` decorator will enable step function mode, which will collect additional meta data about your step functions.
+
 ## Plugins
 
 IOpipe's functionality can be extended through plugins. Plugins hook into the agent lifecycle to allow you to perform additional analytics.