Throw exception when calling response.success()/.failure() if with-block has not been entered

locust/clients.py

@@ -9,7 +9,7 @@
 
 from urllib.parse import urlparse, urlunparse
 
-from .exception import CatchResponseError, ResponseError
+from .exception import CatchResponseError, LocustError, ResponseError
 
 absolute_http_url_regexp = re.compile(r"^https?://", re.I)
 
@@ -216,6 +216,7 @@ class ResponseContextManager(LocustResponse):
     """
 
     _manual_result = None
+    _entered = False
 
     def __init__(self, response, request_event, request_meta):
         # copy data from response to this object
@@ -224,6 +225,7 @@ def __init__(self, response, request_event, request_meta):
         self.request_meta = request_meta
 
     def __enter__(self):
+        self._entered = True
         return self
 
     def __exit__(self, exc, value, traceback):
@@ -267,6 +269,10 @@ def success(self):
                 if response.status_code == 404:
                     response.success()
         """
+        if not self._entered:
+            raise LocustError(
+                "Tried to set status on a request that has not yet been made. Make sure you use a with-block, like this:\n\nwith self.client.request(..., catch_response=True) as response:\n    response.success()"
+            )
         self._manual_result = True
 
     def failure(self, exc):
@@ -282,6 +288,10 @@ def failure(self, exc):
                 if response.content == b"":
                     response.failure("No data")
         """
+        if not self._entered:
+            raise LocustError(
+                "Tried to set status on a request that has not yet been made. Make sure you use a with-block, like this:\n\nwith self.client.request(..., catch_response=True) as response:\n    response.failure(...)"
+            )
         if not isinstance(exc, Exception):
             exc = CatchResponseError(exc)
         self._manual_result = exc

locust/contrib/fasthttp.py

@@ -413,6 +413,7 @@ class ResponseContextManager(FastResponse):
     """
 
     _manual_result = None
+    _entered = False
 
     def __init__(self, response, environment, request_meta):
         # copy data from response to this object
@@ -423,6 +424,7 @@ def __init__(self, response, environment, request_meta):
         self.request_meta = request_meta
 
     def __enter__(self):
+        self._entered = True
         return self
 
     def __exit__(self, exc, value, traceback):
@@ -465,6 +467,10 @@ def success(self):
                 if response.status_code == 404:
                     response.success()
         """
+        if not self._entered:
+            raise LocustError(
+                "Tried to set status on a request that has not yet been made. Make sure you use a with-block, like this:\n\nwith self.client.request(..., catch_response=True) as response:\n    response.success()"
+            )
         self._manual_result = True
 
     def failure(self, exc):
@@ -480,6 +486,10 @@ def failure(self, exc):
                 if response.content == "":
                     response.failure("No data")
         """
+        if not self._entered:
+            raise LocustError(
+                "Tried to set status on a request that has not yet been made. Make sure you use a with-block, like this:\n\nwith self.client.request(..., catch_response=True) as response:\n    response.failure(...)"
+            )
         if not isinstance(exc, Exception):
             exc = CatchResponseError(exc)
         self._manual_result = exc

locust/test/test_fasthttp.py

@@ -6,7 +6,7 @@
 from locust.user import task, TaskSet
 from locust.contrib.fasthttp import FastHttpSession
 from locust import FastHttpUser
-from locust.exception import CatchResponseError, InterruptTaskSet, ResponseError
+from locust.exception import CatchResponseError, InterruptTaskSet, LocustError, ResponseError
 from locust.main import is_user_class
 from .testcases import WebserverTestCase, LocustTestCase
 from .util import create_tls_cert
@@ -459,7 +459,7 @@ def setUp(self):
         class MyUser(FastHttpUser):
             host = "http://127.0.0.1:%i" % self.port
 
-        self.locust = MyUser(self.environment)
+        self.user = MyUser(self.environment)
 
         self.num_failures = 0
         self.num_success = 0
@@ -474,30 +474,30 @@ def on_request(exception, **kwargs):
         self.environment.events.request.add_listener(on_request)
 
     def test_catch_response(self):
-        self.assertEqual(500, self.locust.client.get("/fail").status_code)
+        self.assertEqual(500, self.user.client.get("/fail").status_code)
         self.assertEqual(1, self.num_failures)
         self.assertEqual(0, self.num_success)
 
-        with self.locust.client.get("/ultra_fast", catch_response=True) as response:
+        with self.user.client.get("/ultra_fast", catch_response=True) as response:
             pass
         self.assertEqual(1, self.num_failures)
         self.assertEqual(1, self.num_success)
         self.assertIn("ultra fast", str(response.content))
 
-        with self.locust.client.get("/ultra_fast", catch_response=True) as response:
+        with self.user.client.get("/ultra_fast", catch_response=True) as response:
             raise ResponseError("Not working")
 
         self.assertEqual(2, self.num_failures)
         self.assertEqual(1, self.num_success)
 
     def test_catch_response_http_fail(self):
-        with self.locust.client.get("/fail", catch_response=True) as response:
+        with self.user.client.get("/fail", catch_response=True) as response:
             pass
         self.assertEqual(1, self.num_failures)
         self.assertEqual(0, self.num_success)
 
     def test_catch_response_http_manual_fail(self):
-        with self.locust.client.get("/ultra_fast", catch_response=True) as response:
+        with self.user.client.get("/ultra_fast", catch_response=True) as response:
             response.failure("Haha!")
         self.assertEqual(1, self.num_failures)
         self.assertEqual(0, self.num_success)
@@ -507,13 +507,13 @@ def test_catch_response_http_manual_fail(self):
         )
 
     def test_catch_response_http_manual_success(self):
-        with self.locust.client.get("/fail", catch_response=True) as response:
+        with self.user.client.get("/fail", catch_response=True) as response:
             response.success()
         self.assertEqual(0, self.num_failures)
         self.assertEqual(1, self.num_success)
 
     def test_catch_response_allow_404(self):
-        with self.locust.client.get("/does/not/exist", catch_response=True) as response:
+        with self.user.client.get("/does/not/exist", catch_response=True) as response:
             self.assertEqual(404, response.status_code)
             if response.status_code == 404:
                 response.success()
@@ -561,6 +561,12 @@ class MyUser(FastHttpUser):
         self.assertEqual(0, self.num_success)
         self.assertEqual(1, self.num_failures)
 
+    def test_catch_response_missing_with_block(self):
+        # incorrect usage, missing with-block
+        r = self.user.client.get("/fail", catch_response=True)
+        self.assertRaises(LocustError, r.success)
+        self.assertRaises(LocustError, r.failure, "")
+
     def test_deprecated_request_events(self):
         status = {"success_amount": 0, "failure_amount": 0}
 
@@ -572,9 +578,9 @@ def on_failure(**kw):
 
         self.environment.events.request_success.add_listener(on_success)
         self.environment.events.request_failure.add_listener(on_failure)
-        with self.locust.client.get("/ultra_fast", catch_response=True) as response:
+        with self.user.client.get("/ultra_fast", catch_response=True) as response:
             pass
-        with self.locust.client.get("/wrong_url", catch_response=True) as response:
+        with self.user.client.get("/wrong_url", catch_response=True) as response:
             pass
 
         self.assertEqual(1, status["success_amount"])

locust/test/test_http.py

@@ -4,7 +4,7 @@
 from requests.exceptions import InvalidSchema, InvalidURL, MissingSchema, RequestException
 
 from locust.clients import HttpSession
-from locust.exception import ResponseError
+from locust.exception import LocustError, ResponseError
 from .testcases import WebserverTestCase
 
 
@@ -266,6 +266,13 @@ 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_missing_with_block(self):
+        s = self.get_client()
+        # incorrect usage, missing with-block
+        r = s.get("/fail", catch_response=True)
+        self.assertRaises(LocustError, r.success)
+        self.assertRaises(LocustError, r.failure, "")
+
     def test_user_context(self):
         class TestUser(HttpUser):
             host = f"http://127.0.0.1:{self.port}"