Merge pull request #2511 from locustio/add-event.measure-context-mana…

…ger-for-simpler-firing-of-request-event

Add event.measure context manager for simpler firing of request event

locust/event.py

@@ -1,6 +1,9 @@
 import logging
 from . import log
 import traceback
+from contextlib import contextmanager
+from typing import Generator, Any, Dict
+import time
 from .exception import StopUser, RescheduleTask, RescheduleTaskImmediately, InterruptTaskSet
 
 
@@ -46,6 +49,34 @@ def fire(self, *, reverse=False, **kwargs):
                 logging.error("Uncaught exception in event handler: \n%s", traceback.format_exc())
                 log.unhandled_greenlet_exception = True
 
+    @contextmanager
+    def measure(
+        self, request_type: str, name: str, response_length: int = 0, context=None
+    ) -> Generator[Dict[str, Any], None, None]:
+        """Convenience method for firing the event with automatically calculated response time and automatically marking the request as failed if an exception is raised (this is really only useful for the *request* event)
+
+        Other attributes in the event can be added/updated inside your with block if you use the "with ... as request_meta:" syntax.
+
+        Experimental.
+        """
+        start_time = time.time()
+        start_perf_counter = time.perf_counter()
+        request_meta = {
+            "request_type": request_type,
+            "name": name,
+            "response_length": response_length,
+            "context": context or {},
+            "exception": None,
+            "start_time": start_time,
+        }
+        try:
+            yield request_meta
+        except Exception as e:
+            request_meta["exception"] = e
+        finally:
+            request_meta["response_time"] = (time.perf_counter() - start_perf_counter) * 1000
+            self.fire(**request_meta)
+
 
 class DeprecatedEventHook(EventHook):
     def __init__(self, message):

locust/test/test_http.py

@@ -278,6 +278,29 @@ def test_missing_catch_response_true(self):
         with s.get("/fail") as resp:
             self.assertRaises(LocustError, resp.success)
 
+    def test_event_measure(self):
+        kwargs = {}
+
+        def on_request(**kw):
+            kwargs.update(**kw)
+
+        self.environment.events.request.add_listener(on_request)
+
+        with self.environment.events.request.measure("GET", "/test") as request_meta:
+            time.sleep(0.001)
+
+        self.assertTrue(1 <= kwargs["response_time"] <= 1.5, kwargs["response_time"])
+        self.assertEqual(kwargs["name"], "/test")
+        self.assertIsNone(kwargs["exception"])
+
+        with self.environment.events.request.measure("GET", "/test") as request_meta:
+            request_meta["foo"] = "bar"
+            raise Exception("nooo")
+
+        self.assertEqual(kwargs["name"], "/test")
+        self.assertEqual(kwargs["foo"], "bar")
+        self.assertEqual(str(kwargs["exception"]), "nooo")
+
     def test_user_context(self):
         class TestUser(HttpUser):
             host = f"http://127.0.0.1:{self.port}"