Not closing segment in after_request in case of internal error (#271)

Author: srprash

aws_xray_sdk/ext/flask/middleware.py

@@ -81,6 +81,9 @@ def _after_request(self, response):
         if cont_len:
             segment.put_http_meta(http.CONTENT_LENGTH, int(cont_len))
 
+        if response.status_code >= 500:
+            return response
+
         if self.in_lambda_ctx:
             self._recorder.end_subsegment()
         else:

tests/ext/flask/test_flask.py

@@ -39,8 +39,11 @@ def template():
 recorder.configure(service='test', sampling=False, context=Context())
 XRayMiddleware(app, recorder)
 
-# enable testing mode
-app.config['TESTING'] = True
+# We don't need to enable testing mode by doing app.config['TESTING'] = True
+# because what it does is disable error catching during request handling,
+# so that you get better error reports when performing test requests against the application.
+# But this also results in `after_request` method not getting invoked during unhandled exception which we want
+# since it is the actual application behavior in our use case.
 app = app.test_client()
 
 BASE_URL = 'http://localhost{}'