BugFix: Always close segment in teardown_request handler (aws#272)

* always closing segment if present in teardown_request handler

* close segment in teardown_request only. rename teardown method

Author: srprash

aws_xray_sdk/ext/flask/middleware.py

@@ -17,7 +17,7 @@ def __init__(self, app, recorder):
         self._recorder = recorder
         self.app.before_request(self._before_request)
         self.app.after_request(self._after_request)
-        self.app.teardown_request(self._handle_exception)
+        self.app.teardown_request(self._teardown_request)
         self.in_lambda_ctx = False
 
         if check_in_lambda() and type(self._recorder.context) == LambdaContext:
@@ -81,18 +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:
-            self._recorder.end_segment()
         return response
 
-    def _handle_exception(self, exception):
-        if not exception:
-            return
+    def _teardown_request(self, exception):
         segment = None
         try:
             if self.in_lambda_ctx:
@@ -104,9 +95,11 @@ def _handle_exception(self, exception):
         if not segment:
             return
 
-        segment.put_http_meta(http.STATUS, 500)
-        stack = stacktrace.get_stacktrace(limit=self._recorder._max_trace_back)
-        segment.add_exception(exception, stack)
+        if exception:
+            segment.put_http_meta(http.STATUS, 500)
+            stack = stacktrace.get_stacktrace(limit=self._recorder._max_trace_back)
+            segment.add_exception(exception, stack)
+
         if self.in_lambda_ctx:
             self._recorder.end_subsegment()
         else:

tests/ext/flask/test_flask.py

@@ -9,7 +9,6 @@
 from tests.util import get_new_stubbed_recorder
 import os
 
-
 # define a flask app for testing purpose
 app = Flask(__name__)
 
@@ -29,6 +28,11 @@ def fault():
     return {}['key']
 
 
+@app.route('/fault_no_exception')
+def fault_no_exception():
+    return "SomeException", 500
+
+
 @app.route('/template')
 def template():
     return render_template_string('hello template')
@@ -108,6 +112,18 @@ def test_fault():
     assert exception.type == 'KeyError'
 
 
+def test_fault_no_exception():
+    path = '/fault_no_exception'
+    app.get(path)
+    segment = recorder.emitter.pop()
+    assert not segment.in_progress
+    assert segment.fault
+
+    response = segment.http['response']
+    assert response['status'] == 500
+    assert segment.cause == {}
+
+
 def test_render_template():
     path = '/template'
     app.get(path)