Fix UnboundLocalError when aiohttp server raises a CancelledError (#356)

Co-authored-by: Carol Abadeer <[email protected]>
Co-authored-by: Prashant Srivastava <[email protected]>

Author: 3 people

Diff for: aws_xray_sdk/ext/aiohttp/middleware.py

@@ -67,7 +67,7 @@ async def middleware(request, handler):
         # Non 2XX responses are raised as HTTPExceptions
         response = exc
         six.raise_from(exc, exc)
-    except Exception as exc:
+    except BaseException as exc:
         # Store exception information including the stacktrace to the segment
         response = None
         segment.put_http_meta(http.STATUS, 500)

Diff for: tests/ext/aiohttp/test_middleware.py

@@ -75,9 +75,10 @@ async def handle_unauthorized(self, request: web.Request) -> web.Response:
 
     async def handle_exception(self, request: web.Request) -> web.Response:
         """
-        Handle /exception which raises a KeyError
+        Handle /exception which raises a CancelledError; this is important, as starting from python 3.8 CancelledError
+        extends BaseException instead of Exception
         """
-        return {}['key']
+        raise asyncio.CancelledError()
 
     async def handle_delay(self, request: web.Request) -> web.Response:
         """
@@ -213,8 +214,8 @@ async def test_exception(test_client, loop, recorder):
     """
     client = await test_client(ServerTest.app(loop=loop))
 
-    resp = await client.get('/exception')
-    await resp.text()  # Need this to trigger Exception
+    with pytest.raises(Exception):
+        await client.get('/exception')
 
     segment = recorder.emitter.pop()
     assert not segment.in_progress
@@ -227,7 +228,7 @@ async def test_exception(test_client, loop, recorder):
     assert request['url'] == 'http://127.0.0.1:{port}/exception'.format(port=client.port)
     assert request['client_ip'] == '127.0.0.1'
     assert response['status'] == 500
-    assert exception.type == 'KeyError'
+    assert exception.type == 'CancelledError'
 
 
 async def test_unhauthorized(test_client, loop, recorder):