Fastapi status code ep 6070

epsagon/runners/fastapi.py

@@ -109,6 +109,14 @@ def _update_raw_response_body(self, response, response_type):
         )
 
 
+    def _update_status_code(self, status_code):
+        """
+        Updates the event with given status code.
+        """
+        self.resource['metadata']['status_code'] = status_code
+        if status_code and status_code >= 500:
+            self.set_error()
+
     def _update_raw_response(self, response):
         """
         Updates the event with data by given raw response.
@@ -126,12 +134,9 @@ def _update_raw_response(self, response):
                 'Response Headers',
                 response_headers
             )
+        self._update_status_code(response.status_code)
 
-        self.resource['metadata']['status_code'] = response.status_code
-        if response.status_code >= 500:
-            self.set_error()
-
-    def update_response(self, response):
+    def update_response(self, response, status_code=None):
         """
         Adds response data to event.
         """
@@ -148,3 +153,5 @@ def update_response(self, response):
                 print_debug(
                     'Could not json encode fastapi handler response data'
                 )
+            if status_code:
+                self._update_status_code(status_code)

epsagon/wrappers/fastapi.py

@@ -37,10 +37,11 @@ def _handle_wrapper_params(_args, kwargs, original_request_param_name):
     return kwargs.pop(EPSAGON_REQUEST_PARAM_NAME, None)
 
 
-def _handle_response(response, trace, raised_err):
+def _handle_response(response, status_code, trace, raised_err):
     """
-    Handles the HTTP handler response. Used by both async & sync wrappers.
+    Handles the HTTP handler response.
     :param response:
+    :param status_code:
     :param trace: to update the response with
     :param raised_err: any error which might occured while trying to get
     the response.
@@ -56,7 +57,7 @@ def _handle_response(response, trace, raised_err):
                 return response
 
         if response is not None:
-            trace.runner.update_response(response)
+            trace.runner.update_response(response, status_code=status_code)
         epsagon.trace.trace_factory.send_traces()
     except Exception as exception:  # pylint: disable=W0703
         print_debug('Failed to send traces: {}'.format(exception))
@@ -68,7 +69,7 @@ def _handle_response(response, trace, raised_err):
     return response
 
 # pylint: disable=too-many-statements
-def _wrap_handler(dependant):
+def _wrap_handler(dependant, status_code):
     """
     Wraps the endppint handler.
     """
@@ -150,7 +151,7 @@ def wrapped_handler(*args, **kwargs):
             )
         finally:
             loop.close()
-        return _handle_response(response, trace, raised_err)
+        return _handle_response(response, status_code, trace, raised_err)
 
     dependant.call = wrapped_handler
 
@@ -166,4 +167,4 @@ def __init__(self, *args, **kwargs):
         """
         super().__init__(*args, **kwargs)
         if self.dependant and self.dependant.call:
-            _wrap_handler(self.dependant)
+            _wrap_handler(self.dependant, kwargs.pop('status_code', 200))

tests/wrappers/test_fastapi_wrapper.py

@@ -28,6 +28,10 @@
 CUSTOM_RESPONSE = ["A"]
 CUSTOM_RESPONSE_PATH = "/custom_response"
 BASE_MODEL_RESPONSE_PATH = "/base_model_response"
+DEFAULT_STATUS_CODE = 200
+CUSTOM_STATUS_CODE = 202
+CUSTOM_STATUS_CODE_PATH = "/custom_status_code"
+OVERRIDDEN_CUSTOM_STATUS_CODE_PATH = "/overridden_custom_status_code"
 
 class CustomBaseModel(BaseModel):
     data: List[str]
@@ -48,6 +52,12 @@ def handle_custom_response(response_model=List[str]):
 def handle_base_model_response(response_model=CustomBaseModel):
     return CustomBaseModel(data=CUSTOM_RESPONSE)
 
+def handle_custom_status_code(response_model=List[str]):
+    return CUSTOM_RESPONSE
+
+def handle_overridden_custom_status_code():
+    return _get_response(RETURN_VALUE)
+
 def handle_given_request(request: Request):
     assert request.method == 'POST'
     loop = None
@@ -95,16 +105,29 @@ def fastapi_app():
         handle_base_model_response,
         methods=["GET"]
     )
+    app.add_api_route(
+        CUSTOM_STATUS_CODE_PATH,
+        handle_custom_status_code,
+        methods=["GET"],
+        status_code=CUSTOM_STATUS_CODE
+    )
+    app.add_api_route(
+        OVERRIDDEN_CUSTOM_STATUS_CODE_PATH,
+        handle_overridden_custom_status_code,
+        methods=["GET"],
+        status_code=CUSTOM_STATUS_CODE
+    )
     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"])
-    app.add_api_route("/err", handle_error, methods=["GET"])
+    app.add_api_route("/err", handle_error, methods=["GET"], status_code=200)
     app.add_api_route(MULTIPLE_THREADS_ROUTE, multiple_threads_route, methods=["GET"])
     router = APIRouter()
     router.add_api_route(TEST_ROUTER_PATH, handle_router_endpoint)
     app.include_router(router, prefix=TEST_ROUTER_PREFIX)
     return app
 
+
 @pytest.mark.asyncio
 async def test_fastapi_sanity(trace_transport, fastapi_app):
     """Sanity test."""
@@ -115,6 +138,7 @@ async def test_fastapi_sanity(trace_transport, fastapi_app):
     assert isinstance(runner, FastapiRunner)
     assert runner.resource['name'].startswith('127.0.0.1')
     assert runner.resource['metadata']['Path'] == '/'
+    assert runner.resource['metadata']['status_code'] == DEFAULT_STATUS_CODE
     expected_response_data = _get_response_data(RETURN_VALUE)
     assert runner.resource['metadata']['Response Data'] == (
         expected_response_data
@@ -127,7 +151,7 @@ async def test_fastapi_sanity(trace_transport, fastapi_app):
 
 @pytest.mark.asyncio
 async def test_fastapi_custom_response(trace_transport, fastapi_app):
-    """Sanity test."""
+    """custom response test."""
     request_path = f'{CUSTOM_RESPONSE_PATH}?x=testval'
     async with AsyncClient(app=fastapi_app, base_url="http://test") as ac:
         response = await ac.get(request_path)
@@ -136,6 +160,7 @@ async def test_fastapi_custom_response(trace_transport, fastapi_app):
     assert isinstance(runner, FastapiRunner)
     assert runner.resource['name'].startswith('127.0.0.1')
     assert runner.resource['metadata']['Path'] == CUSTOM_RESPONSE_PATH
+    assert runner.resource['metadata']['status_code'] == DEFAULT_STATUS_CODE
     assert runner.resource['metadata']['Query Params'] == { 'x': 'testval'}
     assert runner.resource['metadata']['Response Data'] == (
         jsonable_encoder(CUSTOM_RESPONSE)
@@ -147,7 +172,7 @@ async def test_fastapi_custom_response(trace_transport, fastapi_app):
 
 @pytest.mark.asyncio
 async def test_fastapi_base_model_response(trace_transport, fastapi_app):
-    """Sanity test."""
+    """base model response 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)
@@ -157,6 +182,7 @@ async def test_fastapi_base_model_response(trace_transport, fastapi_app):
     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']['status_code'] == DEFAULT_STATUS_CODE
     assert runner.resource['metadata']['Query Params'] == { 'x': 'testval'}
     assert runner.resource['metadata']['Response Data'] == (
         jsonable_encoder(expected_response_data)
@@ -166,9 +192,58 @@ async def test_fastapi_base_model_response(trace_transport, fastapi_app):
     assert not trace_factory.traces
 
 
+@pytest.mark.asyncio
+async def test_fastapi_custom_status_code(trace_transport, fastapi_app):
+    """custom status code test."""
+    request_path = f'{CUSTOM_STATUS_CODE_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]
+    assert isinstance(runner, FastapiRunner)
+    assert runner.resource['name'].startswith('127.0.0.1')
+    assert runner.resource['metadata']['Path'] == CUSTOM_STATUS_CODE_PATH
+    assert runner.resource['metadata']['status_code'] == CUSTOM_STATUS_CODE
+    assert runner.resource['metadata']['Query Params'] == { 'x': 'testval'}
+    assert runner.resource['metadata']['Response Data'] == (
+        jsonable_encoder(CUSTOM_RESPONSE)
+    )
+    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_custom_status_code_overridden(trace_transport, fastapi_app):
+    """custom status code test - status code overridden by returned Response """
+    path = OVERRIDDEN_CUSTOM_STATUS_CODE_PATH
+    request_path = f'{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]
+    assert isinstance(runner, FastapiRunner)
+    assert runner.resource['name'].startswith('127.0.0.1')
+    assert runner.resource['metadata']['Path'] == path
+    assert runner.resource['metadata']['status_code'] == DEFAULT_STATUS_CODE
+    assert runner.resource['metadata']['Query Params'] == { 'x': 'testval'}
+    expected_response_data = _get_response_data(RETURN_VALUE)
+    assert runner.resource['metadata']['Response Data'] == (
+        expected_response_data
+    )
+    assert runner.resource['metadata']['Query Params'] == { 'x': 'testval'}
+    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."""
+    """handler with a request parameter test."""
     request_path = f'{REQUEST_OBJ_PATH}?x=testval'
     async with AsyncClient(app=fastapi_app, base_url="http://test") as ac:
         response = await ac.post(request_path, json=TEST_POST_DATA)
@@ -177,6 +252,7 @@ async def test_fastapi_given_request(trace_transport, fastapi_app):
     assert isinstance(runner, FastapiRunner)
     assert runner.resource['name'].startswith('127.0.0.1')
     assert runner.resource['metadata']['Path'] == REQUEST_OBJ_PATH
+    assert runner.resource['metadata']['status_code'] == DEFAULT_STATUS_CODE
     expected_response_data = _get_response_data(RETURN_VALUE)
     assert runner.resource['metadata']['Response Data'] == (
         expected_response_data
@@ -198,6 +274,7 @@ async def test_fastapi_custom_router(trace_transport, fastapi_app):
     assert isinstance(runner, FastapiRunner)
     assert runner.resource['name'].startswith('127.0.0.1')
     assert runner.resource['metadata']['Path'] == full_route_path
+    assert runner.resource['metadata']['status_code'] == DEFAULT_STATUS_CODE
     expected_response_data = _get_response_data(ROUTER_RETURN_VALUE)
     assert runner.resource['metadata']['Response Data'] == (
         expected_response_data
@@ -218,6 +295,7 @@ async def test_fastapi_exception(trace_transport, fastapi_app):
     assert runner.error_code == ErrorCode.EXCEPTION
     assert runner.exception['type'] == 'CustomFastAPIException'
     assert runner.exception['message'] == 'test'
+    assert "status_code" not in runner.resource['metadata']
     # validating no `zombie` traces exist
     assert not trace_factory.traces
 
@@ -232,6 +310,7 @@ async def _send_request(app, path, trace_transport):
     assert isinstance(runner, FastapiRunner)
     assert runner.resource['name'].startswith('127.0.0.1')
     assert runner.resource['metadata']['Path'] == request_path
+    assert runner.resource['metadata']['status_code'] == DEFAULT_STATUS_CODE
     expected_response_data = _get_response_data(path)
     assert runner.resource['metadata']['Response Data'] == (
         expected_response_data
@@ -280,6 +359,7 @@ def test_fastapi_multiple_threads_route(trace_transport, fastapi_app):
     assert isinstance(runner, FastapiRunner)
     assert runner.resource['name'].startswith('127.0.0.1')
     assert runner.resource['metadata']['Path'] == MULTIPLE_THREADS_ROUTE
+    assert runner.resource['metadata']['status_code'] == DEFAULT_STATUS_CODE
     expected_response_data = _get_response_data(MULTIPLE_THREADS_RETURN_VALUE)
     assert runner.resource['metadata']['Response Data'] == (
         expected_response_data