Merge pull request #2276 from ianmetcalf/response-context-manager

Fix exception grouping for requests with both catch_response and name arguments

locust/clients.py

@@ -135,8 +135,11 @@ def request(self, method, url, name=None, catch_response=False, context={}, **kw
         response = self._send_request_safe_mode(method, url, **kwargs)
         response_time = (time.perf_counter() - start_perf_counter) * 1000
 
-        request_after_redirect = (response.history and response.history[0] or response).request
-        url_after_redirect = request_after_redirect.path_url
+        request_before_redirect = (response.history and response.history[0] or response).request
+        url = request_before_redirect.url
+
+        if not name:
+            name = request_before_redirect.path_url
 
         if self.user:
             context = {**self.user.context(), **context}
@@ -145,12 +148,12 @@ def request(self, method, url, name=None, catch_response=False, context={}, **kw
         request_meta = {
             "request_type": method,
             "response_time": response_time,
-            "name": name or url_after_redirect,
+            "name": name,
             "context": context,
             "response": response,
             "exception": None,
             "start_time": start_time,
-            "url": request_after_redirect.url,
+            "url": url,
         }
 
         # get the length of the content, but if the argument stream is set to True, we take
@@ -163,34 +166,8 @@ def request(self, method, url, name=None, catch_response=False, context={}, **kw
         if catch_response:
             return ResponseContextManager(response, request_event=self.request_event, request_meta=request_meta)
         else:
-            if name:
-                # Since we use the Exception message when grouping failures, in order to not get
-                # multiple failure entries for different URLs for the same name argument, we need
-                # to temporarily override the response.url attribute
-                orig_url = response.url
-                response.url = name
-
-            try:
-                response.raise_for_status()
-            except RequestException as e:
-                while (
-                    isinstance(
-                        e,
-                        (
-                            requests.exceptions.ConnectionError,
-                            requests.packages.urllib3.exceptions.ProtocolError,
-                            requests.packages.urllib3.exceptions.MaxRetryError,
-                            requests.packages.urllib3.exceptions.NewConnectionError,
-                        ),
-                    )
-                    and e.__context__  # Not sure if the above exceptions can ever be the lowest level, but it is good to be sure
-                ):
-                    e = e.__context__
-                request_meta["exception"] = e
-
-            self.request_event.fire(**request_meta)
-            if name:
-                response.url = orig_url
+            with ResponseContextManager(response, request_event=self.request_event, request_meta=request_meta):
+                pass
             return response
 
     def _send_request_safe_mode(self, method, url, **kwargs):
@@ -253,12 +230,32 @@ def __exit__(self, exc, value, traceback):
                 # we want other unknown exceptions to be raised
                 return False
         else:
+            # Since we use the Exception message when grouping failures, in order to not get
+            # multiple failure entries for different URLs for the same name argument, we need
+            # to temporarily override the response.url attribute
+            orig_url = self.url
+            self.url = self.request_meta["name"]
+
             try:
                 self.raise_for_status()
             except requests.exceptions.RequestException as e:
+                while (
+                    isinstance(
+                        e,
+                        (
+                            requests.exceptions.ConnectionError,
+                            requests.packages.urllib3.exceptions.ProtocolError,
+                            requests.packages.urllib3.exceptions.MaxRetryError,
+                            requests.packages.urllib3.exceptions.NewConnectionError,
+                        ),
+                    )
+                    and e.__context__  # Not sure if the above exceptions can ever be the lowest level, but it is good to be sure
+                ):
+                    e = e.__context__
                 self.request_meta["exception"] = e
 
             self._report_request()
+            self.url = orig_url
 
         return True
 

locust/test/test_http.py

@@ -266,6 +266,22 @@ def test_catch_response_default_fail(self):
         self.assertEqual(1, self.environment.stats.total.num_requests)
         self.assertEqual(1, self.environment.stats.total.num_failures)
 
+    def test_catch_response_with_name_replacement(self):
+        s = self.get_client()
+        kwargs = {}
+
+        def on_request(**kw):
+            self.assertIsNotNone(kw["exception"])
+            kwargs.update(kw)
+
+        self.environment.events.request.add_listener(on_request)
+
+        with s.get("/wrong_url/01", name="replaced_url_name") as r:
+            pass
+
+        self.assertIn("for url: replaced_url_name", str(kwargs["exception"]))
+        self.assertEqual(s.base_url + "/wrong_url/01", kwargs["url"])  # url is unaffected by name
+
     def test_catch_response_missing_with_block(self):
         s = self.get_client()
         # incorrect usage, missing with-block