Update things for fastapi tests and add tests

rollbar/__init__.py

@@ -84,6 +84,11 @@
 except ImportError:
     SanicRequest = None
 
+try:
+    from fastapi import Request as FastapiRequest
+except ImportError:
+    FastapiRequest = None
+
 try:
     from google.appengine.api.urlfetch import fetch as AppEngineFetch
 except ImportError:
@@ -1117,6 +1122,10 @@ def _build_request_data(request):
     if FalconRequest and isinstance(request, FalconRequest):
         return _build_falcon_request_data(request)
 
+    # fastapi
+    if FastapiRequest and isinstance(request, FastapiRequest):
+        return _build_fastapi_request_data(request)
+
     # Plain wsgi (should be last)
     if isinstance(request, dict) and 'wsgi.version' in request:
         return _build_wsgi_request_data(request)
@@ -1187,6 +1196,24 @@ def _build_django_request_data(request):
     return request_data
 
 
+def _build_fastapi_request_data(request):
+    """Fastapi relies on starlette which uses async functions to retrieve the request data and body.
+
+    So skipping for now since function callers are synchronous...
+    """
+    request_data = {
+        'url': str(request.url),
+        'method': request.method,
+        'GET': None,
+        'POST': None,
+        'user_ip': _asgi_extract_user_ip(request)
+    }
+
+    request_data['headers'] = request.headers
+
+    return request_data
+
+
 def _build_werkzeug_request_data(request):
     request_data = {
         'url': request.url,
@@ -1623,3 +1650,13 @@ def _wsgi_extract_user_ip(environ):
     if real_ip:
         return real_ip
     return environ['REMOTE_ADDR']
+
+
+def _asgi_extract_user_ip(request):
+    forwarded_for = request.headers.get('x-forwarded-for')
+    if forwarded_for:
+        return forwarded_for
+    real_ip = request.headers.get('x-real-ip')
+    if real_ip:
+        return real_ip
+    return request.client.host

rollbar/test/fastapi_tests/test_fastapi.py

@@ -0,0 +1,188 @@
+"""
+Tests for fastapi instrumentation
+"""
+
+import json
+import os
+import sys
+
+try:
+    from unittest import mock
+except ImportError:
+    import mock
+
+import rollbar
+from rollbar.test import BaseTest
+
+# access token for https://rollbar.com/rollbar/pyrollbar
+TOKEN = "92c10f5616944b81a2e6f3c6493a0ec2"
+
+# Fastapi works on python +3.6
+ALLOWED_PYTHON_VERSION = sys.version_info[0] == 3 and sys.version_info[1] >= 6
+
+
+try:
+    import fastapi  # noqa: F401
+
+    FASTAPI_INSTALLED = True
+except ImportError:
+    FASTAPI_INSTALLED = False
+
+
+def create_app():
+    from fastapi import FastAPI
+
+    app = FastAPI()
+
+    @app.get("/")
+    def index():
+        return "Index page"
+
+    @app.get("/cause_error")
+    @app.post("/cause_error")
+    def cause_error():
+        raise Exception("Uh oh")
+
+    return app
+
+
+def init_rollbar(app):
+    import rollbar.contrib.fastapi
+    from fastapi import Request, status
+    from fastapi.responses import Response
+
+    rollbar.init(
+        TOKEN,
+        "fastapitest",
+        root=os.path.dirname(os.path.realpath(__file__)),
+        allow_logging_basic_config=True,
+        capture_email=True,
+        capture_username=True,
+    )
+
+    @app.exception_handler(Exception)
+    async def handle_unexpected_exceptions(request: Request, exc: Exception):
+        """This won't capture HTTPException."""
+        try:
+            raise exc
+        except Exception:
+            rollbar.contrib.fastapi.report_exception(request=request)
+
+        return Response(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR)
+
+
+if ALLOWED_PYTHON_VERSION and FASTAPI_INSTALLED:
+
+    from fastapi.testclient import TestClient
+
+    class FastapiTest(BaseTest):
+        def setUp(self):
+            super().setUp()
+            self.app = create_app()
+            init_rollbar(self.app)
+            self.client = TestClient(self.app, raise_server_exceptions=False)
+
+        def test_index(self):
+            resp = self.client.get("/")
+            self.assertEqual(resp.status_code, 200)
+
+        def assertStringEqual(self, left, right):
+            if sys.version_info[0] > 2:
+                if hasattr(left, "decode"):
+                    left = left.decode("ascii")
+                if hasattr(right, "decode"):
+                    right = right.decode("ascii")
+
+                return self.assertEqual(left, right)
+            else:
+                return self.assertEqual(left, right)
+
+        @mock.patch("rollbar.send_payload")
+        def test_uncaught(self, send_payload):
+            resp = self.client.get(
+                "/cause_error?foo=bar",
+                headers={"X-Real-Ip": "1.2.3.4", "User-Agent": "Fastapi Test"},
+            )
+            self.assertEqual(resp.status_code, 500)
+
+            self.assertEqual(send_payload.called, True)
+            payload = send_payload.call_args[0][0]
+            data = payload["data"]
+
+            self.assertIn("body", data)
+            self.assertEqual(data["body"]["trace"]["exception"]["class"], "Exception")
+            self.assertStringEqual(
+                data["body"]["trace"]["exception"]["message"], "Uh oh"
+            )
+
+            self.assertIn("request", data)
+            self.assertEqual(
+                data["request"]["url"], "http://testserver/cause_error?foo=bar"
+            )
+
+            self.assertEqual(data["request"]["user_ip"], "1.2.3.4")
+            self.assertEqual(data["request"]["method"], "GET")
+            self.assertEqual(data["request"]["headers"]["user-agent"], "Fastapi Test")
+
+        @mock.patch("rollbar.send_payload")
+        def test_uncaught_json_request(self, send_payload):
+            json_body = {"hello": "world"}
+            json_body_str = json.dumps(json_body)
+            resp = self.client.post(
+                "/cause_error",
+                data=json_body_str,
+                headers={
+                    "Content-Type": "application/json",
+                    "X-Forwarded-For": "5.6.7.8",
+                },
+            )
+
+            self.assertEqual(resp.status_code, 500)
+
+            self.assertEqual(send_payload.called, True)
+            payload = send_payload.call_args[0][0]
+            data = payload["data"]
+
+            self.assertIn("body", data)
+            self.assertEqual(data["body"]["trace"]["exception"]["class"], "Exception")
+            self.assertStringEqual(
+                data["body"]["trace"]["exception"]["message"], "Uh oh"
+            )
+
+            self.assertIn("request", data)
+            self.assertEqual(data["request"]["url"], "http://testserver/cause_error")
+            self.assertEqual(data["request"]["user_ip"], "5.6.7.8")
+            self.assertEqual(data["request"]["method"], "POST")
+
+        @mock.patch("rollbar.send_payload")
+        def test_uncaught_no_username_no_email(self, send_payload):
+            rollbar.SETTINGS["capture_email"] = False
+            rollbar.SETTINGS["capture_username"] = False
+
+            resp = self.client.get(
+                "/cause_error?foo=bar",
+                headers={"X-Real-Ip": "1.2.3.4", "User-Agent": "Fastapi Test"},
+            )
+            self.assertEqual(resp.status_code, 500)
+
+            self.assertEqual(send_payload.called, True)
+            payload = send_payload.call_args[0][0]
+            data = payload["data"]
+
+            self.assertIn("body", data)
+            self.assertEqual(data["body"]["trace"]["exception"]["class"], "Exception")
+            self.assertStringEqual(
+                data["body"]["trace"]["exception"]["message"], "Uh oh"
+            )
+
+            self.assertIn("request", data)
+            self.assertEqual(
+                data["request"]["url"], "http://testserver/cause_error?foo=bar"
+            )
+
+            self.assertEqual(data["request"]["user_ip"], "1.2.3.4")
+            self.assertEqual(data["request"]["method"], "GET")
+            self.assertEqual(data["request"]["headers"]["user-agent"], "Fastapi Test")
+
+            rollbar.SETTINGS["capture_email"] = True
+            rollbar.SETTINGS["capture_username"] = True