Add fastapi source

Author: tomaisthorpe

aikido_zen/__init__.py

@@ -61,6 +61,7 @@ def protect(mode="daemon", token=""):
     import aikido_zen.sources.flask
     import aikido_zen.sources.quart
     import aikido_zen.sources.starlette
+    import aikido_zen.sources.fastapi
     import aikido_zen.sources.xml_sources.xml
     import aikido_zen.sources.xml_sources.lxml
 

aikido_zen/sources/fastapi.py

@@ -0,0 +1,14 @@
+from aikido_zen.sinks import on_import, patch_function
+from aikido_zen.sources.starlette.starlette_routing import _request_response
+
+
+@on_import("fastapi.routing", "fastapi")
+def patch(m):
+    """
+    FastAPI defines its own request_response in fastapi/routing.py rather than
+    importing it from starlette.routing. Patching only starlette.routing.request_response
+    leaves all FastAPI APIRoute endpoints unwrapped, so pre_response and post_response
+    hooks never fire for those routes.
+    (github src: https://github.com/tiangolo/fastapi/blob/master/fastapi/routing.py)
+    """
+    patch_function(m, "request_response", _request_response)

end2end/fastapi_postgres_uvicorn_test.py

@@ -88,26 +88,30 @@ def test_routes_discovered_in_heartbeat():
 
     validate_heartbeat(
         heartbeat_events[0],
-        routes=[{
-            "apispec": {
-                "body": {
-                    "type": "form-urlencoded",
-                    "schema": {
-                        "type": "object",
-                        "properties": {
-                            "dog_name": {
-                                "items": {"type": "string"},
-                                "type": "array",
-                            }
+        routes=[
+            {
+                "apispec": {
+                    "auth": None,
+                    "body": {
+                        "schema": {
+                            "properties": {"dog_name": {"type": "string"}},
+                            "type": "object",
                         },
+                        "type": "form-urlencoded",
                     },
+                    "query": None,
                 },
-                "query": None,
-                "auth": None,
+                "hits": 2,
+                "hits_delta_since_sync": 0,
+                "method": "POST",
+                "path": "/create",
             },
-            "hits": 1,
-            "hits_delta_since_sync": 0,
-            "method": "POST",
-            "path": "/create",
-        }],
+            {
+                "apispec": {},
+                "hits": 1,
+                "hits_delta_since_sync": 0,
+                "method": "GET",
+                "path": "/sync_route",
+            },
+        ],
     )