fix(fastapi): collecting all responses using fast json encode (#333)

Author: maorlx

epsagon/runners/fastapi.py

@@ -14,12 +14,13 @@
     UJSONResponse,
     RedirectResponse,
 )
+from fastapi.encoders import jsonable_encoder
 from epsagon.common import EpsagonWarning
 from ..event import BaseEvent
-from ..utils import add_data_if_needed, normalize_http_url
+from ..utils import add_data_if_needed, normalize_http_url, print_debug
 from ..constants import EPSAGON_HEADER
 
-SUPPORTED_RESPONSE_TYPES = (
+SUPPORTED_RAW_RESPONSE_TYPES = (
     JSONResponse,
     HTMLResponse,
     PlainTextResponse,
@@ -85,39 +86,65 @@ def update_request_body(self, body):
                 body
             )
 
+    def _update_raw_response_body(self, response, response_type):
+        """
+        Updates the response body by given `raw` response and its type
+        """
+        body = response.body
+        if response_type == JSONResponse:
+            try:
+                body = json.loads(body)
+            except Exception: # pylint: disable=W0703
+                warnings.warn(
+                    'Could not load response json',
+                    EpsagonWarning
+                )
+                body = body.decode('utf-8')
+        else:
+            body = body.decode('utf-8')
+        add_data_if_needed(
+            self.resource['metadata'],
+            'Response Data',
+            body
+        )
+
+
+    def _update_raw_response(self, response):
+        """
+        Updates the event with data by given raw response.
+        Raw response is an instance of Response.
+        """
+        for response_type in SUPPORTED_RAW_RESPONSE_TYPES:
+            if isinstance(response, response_type):
+                self._update_raw_response_body(response, response_type)
+                break
+
+        response_headers = dict(response.headers.items())
+        if response.headers:
+            add_data_if_needed(
+                self.resource['metadata'],
+                'Response Headers',
+                response_headers
+            )
+
+        self.resource['metadata']['status_code'] = response.status_code
+        if response.status_code >= 500:
+            self.set_error()
+
     def update_response(self, response):
         """
         Adds response data to event.
         """
-        for response_type in SUPPORTED_RESPONSE_TYPES:
-            if isinstance(response, response_type):
-                body = response.body
-                if response_type == JSONResponse:
-                    try:
-                        body = json.loads(body)
-                    except Exception: # pylint: disable=W0703
-                        warnings.warn(
-                            'Could not load response json',
-                            EpsagonWarning
-                        )
-                        body = body.decode('utf-8')
-                else:
-                    body = body.decode('utf-8')
+        if isinstance(response, Response):
+            self._update_raw_response(response)
+        else:
+            try:
                 add_data_if_needed(
                     self.resource['metadata'],
                     'Response Data',
-                    body
+                    jsonable_encoder(response)
                 )
-        if isinstance(response, Response):
-            response_headers = dict(response.headers.items())
-            if response.headers:
-                add_data_if_needed(
-                    self.resource['metadata'],
-                    'Response Headers',
-                    response_headers
+            except Exception: # pylint: disable=W0703
+                print_debug(
+                    'Could not json encode fastapi handler response data'
                 )
-
-            self.resource['metadata']['status_code'] = response.status_code
-
-            if response.status_code >= 500:
-                self.set_error()

tests/wrappers/test_fastapi_wrapper.py

@@ -7,8 +7,10 @@
 import asyncio
 from typing import List
 from httpx import AsyncClient
+from pydantic import BaseModel
 from fastapi import FastAPI, APIRouter, Request
 from fastapi.responses import JSONResponse
+from fastapi.encoders import jsonable_encoder
 from epsagon import trace_factory
 from epsagon.common import ErrorCode
 from epsagon.runners.fastapi import FastapiRunner
@@ -25,6 +27,10 @@
 TEST_POST_DATA = {'post_test': '123'}
 CUSTOM_RESPONSE = ["A"]
 CUSTOM_RESPONSE_PATH = "/custom_response"
+BASE_MODEL_RESPONSE_PATH = "/base_model_response"
+
+class CustomBaseModel(BaseModel):
+    data: List[str]
 
 def _get_response_data(key):
     return {key: key}
@@ -39,6 +45,9 @@ def handle():
 def handle_custom_response(response_model=List[str]):
     return CUSTOM_RESPONSE
 
+def handle_base_model_response(response_model=CustomBaseModel):
+    return CustomBaseModel(data=CUSTOM_RESPONSE)
+
 def handle_given_request(request: Request):
     assert request.method == 'POST'
     loop = None
@@ -76,7 +85,16 @@ def handle_error():
 def fastapi_app():
     app = FastAPI()
     app.add_api_route("/", handle, methods=["GET"])
-    app.add_api_route(CUSTOM_RESPONSE_PATH, handle_custom_response, methods=["GET"])
+    app.add_api_route(
+        CUSTOM_RESPONSE_PATH,
+        handle_custom_response,
+        methods=["GET"]
+    )
+    app.add_api_route(
+        BASE_MODEL_RESPONSE_PATH,
+        handle_base_model_response,
+        methods=["GET"]
+    )
     app.add_api_route(REQUEST_OBJ_PATH, handle_given_request, methods=["POST"])
     app.add_api_route("/a", handle_a, methods=["GET"])
     app.add_api_route("/b", handle_b, methods=["GET"])
@@ -119,11 +137,35 @@ async def test_fastapi_custom_response(trace_transport, fastapi_app):
     assert runner.resource['name'].startswith('127.0.0.1')
     assert runner.resource['metadata']['Path'] == CUSTOM_RESPONSE_PATH
     assert runner.resource['metadata']['Query Params'] == { 'x': 'testval'}
+    assert runner.resource['metadata']['Response Data'] == (
+        jsonable_encoder(CUSTOM_RESPONSE)
+    )
     assert response_data == CUSTOM_RESPONSE
     # validating no `zombie` traces exist
     assert not trace_factory.traces
 
 
+@pytest.mark.asyncio
+async def test_fastapi_base_model_response(trace_transport, fastapi_app):
+    """Sanity test."""
+    request_path = f'{BASE_MODEL_RESPONSE_PATH}?x=testval'
+    async with AsyncClient(app=fastapi_app, base_url="http://test") as ac:
+        response = await ac.get(request_path)
+    response_data = response.json()
+    runner = trace_transport.last_trace.events[0]
+    expected_response_data = CustomBaseModel(data=CUSTOM_RESPONSE)
+    assert isinstance(runner, FastapiRunner)
+    assert runner.resource['name'].startswith('127.0.0.1')
+    assert runner.resource['metadata']['Path'] == BASE_MODEL_RESPONSE_PATH
+    assert runner.resource['metadata']['Query Params'] == { 'x': 'testval'}
+    assert runner.resource['metadata']['Response Data'] == (
+        jsonable_encoder(expected_response_data)
+    )
+    assert response_data == expected_response_data
+    # validating no `zombie` traces exist
+    assert not trace_factory.traces
+
+
 @pytest.mark.asyncio
 async def test_fastapi_given_request(trace_transport, fastapi_app):
     """Sanity test."""