1import importlib
2import sys
3
4try:
5    from unittest import mock
6except ImportError:
7    import mock
8
9try:
10    import fastapi
11
12    FASTAPI_INSTALLED = True
13except ImportError:
14    FASTAPI_INSTALLED = False
15
16import unittest2
17
18import rollbar
19from rollbar.test import BaseTest
20
21ALLOWED_PYTHON_VERSION = sys.version_info >= (3, 6)
22
23
24@unittest2.skipUnless(
25    FASTAPI_INSTALLED and ALLOWED_PYTHON_VERSION,
26    'FastAPI LoggerMiddleware requires Python3.6+',
27)
28class LoggerMiddlewareTest(BaseTest):
29    def setUp(self):
30        importlib.reload(rollbar)
31
32    @mock.patch('rollbar._check_config', return_value=True)
33    @mock.patch('rollbar.send_payload')
34    def test_should_add_framework_version_to_payload(self, mock_send_payload, *mocks):
35        import fastapi
36        from fastapi import FastAPI
37        import rollbar
38        from rollbar.contrib.fastapi.logger import LoggerMiddleware
39
40        self.assertIsNone(rollbar.BASE_DATA_HOOK)
41
42        app = FastAPI()
43        app.add_middleware(LoggerMiddleware)
44
45        rollbar.report_exc_info()
46
47        mock_send_payload.assert_called_once()
48        payload = mock_send_payload.call_args[0][0]
49
50        self.assertIn('fastapi', payload['data']['framework'])
51        self.assertIn(fastapi.__version__, payload['data']['framework'])
52
53    def test_should_support_type_hints(self):
54        from starlette.types import Receive, Scope, Send
55        import rollbar.contrib.fastapi.logger
56
57        self.assertDictEqual(
58            rollbar.contrib.fastapi.logger.LoggerMiddleware.__call__.__annotations__,
59            {'scope': Scope, 'receive': Receive, 'send': Send, 'return': None},
60        )
61
62    @mock.patch('rollbar.contrib.starlette.logger.store_current_request')
63    def test_should_store_current_request(self, store_current_request):
64        from fastapi import FastAPI
65        from rollbar.contrib.fastapi.logger import LoggerMiddleware
66
67        try:
68            from fastapi.testclient import TestClient
69        except ImportError:  # Added in FastAPI v0.51.0+
70            from starlette.testclient import TestClient
71
72        expected_scope = {
73            'client': ['testclient', 50000],
74            'headers': [
75                (b'host', b'testserver'),
76                (b'user-agent', b'testclient'),
77                (b'accept-encoding', b'gzip, deflate'),
78                (b'accept', b'*/*'),
79                (b'connection', b'keep-alive'),
80            ],
81            'http_version': '1.1',
82            'method': 'GET',
83            'path': '/',
84            'query_string': b'',
85            'root_path': '',
86            'scheme': 'http',
87            'server': ['testserver', 80],
88            'type': 'http',
89        }
90
91        app = FastAPI()
92        app.add_middleware(LoggerMiddleware)
93
94        @app.get('/')
95        async def read_root():
96            return 'ok'
97
98        client = TestClient(app)
99        client.get('/')
100
101        store_current_request.assert_called_once()
102
103        scope = store_current_request.call_args[0][0]
104        self.assertDictContainsSubset(expected_scope, scope)
105
106    def test_should_return_current_request(self):
107        from fastapi import FastAPI
108        from starlette.requests import Request
109        from rollbar.contrib.fastapi.logger import LoggerMiddleware
110        from rollbar.contrib.fastapi import get_current_request
111
112        try:
113            from fastapi.testclient import TestClient
114        except ImportError:  # Added in FastAPI v0.51.0+
115            from starlette.testclient import TestClient
116
117        app = FastAPI()
118        app.add_middleware(LoggerMiddleware)
119
120        @app.get('/')
121        async def read_root():
122            request = get_current_request()
123
124            self.assertIsNotNone(request)
125            self.assertIsInstance(request, Request)
126
127        client = TestClient(app)
128        client.get('/')
129
130    @mock.patch('rollbar.contrib.starlette.requests.ContextVar', None)
131    @mock.patch('logging.Logger.error')
132    def test_should_not_return_current_request_for_older_python(self, mock_log):
133        from fastapi import FastAPI
134        from rollbar.contrib.fastapi.logger import LoggerMiddleware
135        from rollbar.contrib.fastapi import get_current_request
136
137        try:
138            from fastapi.testclient import TestClient
139        except ImportError:  # Added in FastAPI v0.51.0+
140            from starlette.testclient import TestClient
141
142        app = FastAPI()
143        app.add_middleware(LoggerMiddleware)
144
145        @app.get('/')
146        async def read_root():
147            self.assertIsNone(get_current_request())
148            mock_log.assert_called_once_with(
149                'Python 3.7+ (or aiocontextvars package)'
150                ' is required to receive current request.'
151            )
152
153        client = TestClient(app)
154        client.get('/')
155