fix handling of responses with aiohttp4

elasticapm/contrib/aiohttp/middleware.py

@@ -29,7 +29,7 @@
 #  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 import aiohttp
-from aiohttp.web import Response, middleware
+from aiohttp.web import HTTPException, Response, middleware
 
 import elasticapm
 from elasticapm.conf import constants
@@ -82,13 +82,14 @@ async def handle_request(request, handler):
                     context={"request": get_data_from_request(request, elasticapm_client.config, constants.ERROR)}
                 )
                 elasticapm.set_transaction_result("HTTP 5xx", override=False)
-                if isinstance(exc, Response):
+                elasticapm.set_context({"status_code": 500}, "response")
+                # some exceptions are response-like, e.g. have headers and status code. Let's try and capture them
+                if isinstance(exc, (Response, HTTPException)):
                     elasticapm.set_context(
                         lambda: get_data_from_response(exc, elasticapm_client.config, constants.ERROR),  # noqa: F821
                         "response",
                     )
-                else:
-                    elasticapm.set_context({"status_code": 500}, "response")
+
             raise
         finally:
             elasticapm_client.end_transaction()

elasticapm/contrib/aiohttp/utils.py

@@ -27,8 +27,9 @@
 #  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 #  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 #  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+from typing import Union
 
-from aiohttp.web import Request, Response
+from aiohttp.web import HTTPException, Request, Response
 
 from elasticapm.conf import Config
 from elasticapm.utils import compat, get_url_dict
@@ -49,11 +50,11 @@ def get_data_from_request(request: Request, config: Config, event_type: str):
     return result
 
 
-def get_data_from_response(response: Response, config: Config, event_type: str):
+def get_data_from_response(response: Union[HTTPException, Response], config: Config, event_type: str):
     result = {}
-
-    if isinstance(getattr(response, "status", None), compat.integer_types):
-        result["status_code"] = response.status
+    status = getattr(response, "status", getattr(response, "status_code", None))
+    if isinstance(status, compat.integer_types):
+        result["status_code"] = status
     if config.capture_headers and getattr(response, "headers", None):
         headers = response.headers
         result["headers"] = {key: ";".join(headers.getall(key)) for key in compat.iterkeys(headers)}