Merge pull request #1750 from DennisKrone/feature-1592

Updated request event with context

.git-blame-ignore-revs

@@ -1,2 +1,3 @@
 # Migrate code style to Black
-7c0fcc213d3988f6e7c6ffef63b24afe00e5fbd9
+7c0fcc213d3988f6e7c6ffef63b24afe00e5fbd9
+2e7a8b5697a98d1d314d6fc3ef0589f81f09d7fe

docs/api.rst

@@ -7,7 +7,7 @@ User class
 ============
 
 .. autoclass:: locust.User
-    :members: wait_time, tasks, weight, abstract, on_start, on_stop, wait
+    :members: wait_time, tasks, weight, abstract, on_start, on_stop, wait, context
 
 HttpUser class
 ================

docs/changelog.rst

@@ -4,6 +4,12 @@ Changelog Highlights
 
 For full details of the Locust changelog, please see https://github.com/locustio/locust/blob/master/CHANGELOG.md
 
+1.5.0
+=====
+
+* Add new event called request. Is called on every request successful or not. request_success and request_failure are still available but are deprecated
+* Add parameter context to the request event. Can be used to forward information when calling a request, things like user information, tags etc
+
 1.4.4
 =====
 

docs/extending-locust.rst

@@ -15,9 +15,12 @@ Here's an example on how to set up an event listener::
 
     from locust import events
     
-    @events.request_success.add_listener
-    def my_success_handler(request_type, name, response_time, response_length, **kw):
-        print("Successfully made a request to: %s" % name)
+    @events.request.add_listener
+    def my_request_handler(request_type, name, response_time, response_length, context, exception, **kw):
+        if exception:
+            print(f"Request to {name} failed with exception {exception}")
+        else:
+            print(f"Successfully made a request to: {name})
 
 
 .. note::
@@ -26,12 +29,46 @@ Here's an example on how to set up an event listener::
     (the \**kw in the code above), to prevent your code from breaking if new arguments are
     added in a future version.
 
+
+Request context
+==================
+
+By using the context parameter in the request method information you can attach data that will be forwarded by the 
+request event. This should be a dictionary and can be set directly when calling request() or on a class level 
+by overwriting the User.context() method. 
+
+Context from request method::
+
+    class MyUser(HttpUser):
+        @task
+        def t(self):
+            self.client.post("/login", json={"username": "foo"}, context={"username": "foo"})
+
+        @events.request.add_listener
+        def on_request(context, **kwargs):
+            print(context["username"])
+    
+Context from User class::
+
+    class MyUser(HttpUser):
+        def context(self):
+            return {"username": self.username}
+
+        @task
+        def t(self):
+            self.username = "foo"
+            self.client.post("/login", json={"username": self.username})
+
+        @events.request.add_listener
+        def on_request(self, context, **kwargs):
+            print(context["username"])
+
+
 .. seealso::
 
     To see all available events, please see :ref:`events`.
 
 
-
 Adding Web Routes
 ==================
 

docs/testing-other-systems.rst

@@ -6,8 +6,7 @@ Testing other systems using custom clients
 
 Locust was built with HTTP as its main target. However, it can easily be extended to load test 
 any request/response based system, by writing a custom client that triggers 
-:py:attr:`request_success <locust.event.Events.request_success>` and 
-:py:attr:`request_failure <locust.event.Events.request_failure>` events.
+:py:attr:`request <locust.event.Events.request>`
 
 .. note::
 
@@ -32,8 +31,8 @@ using ``abstract = True`` which means that Locust will not try to create simulat
 
 The ``XmlRpcClient`` is a wrapper around the standard 
 library's :py:class:`xmlrpc.client.ServerProxy`. It basically just proxies the function calls, but with the
-important addition of firing :py:attr:`locust.event.Events.request_success` and :py:attr:`locust.event.Events.request_failure` 
-events, which will record all calls in Locust's statistics.
+important addition of firing :py:attr:`locust.event.Events.request`
+event, which will record all calls in Locust's statistics.
 
 Here's an implementation of an XML-RPC server that would work as a server for the code above:
 

examples/custom_xmlrpc_client/xmlrpc_locustfile.py

@@ -7,8 +7,8 @@
 class XmlRpcClient(ServerProxy):
     """
     Simple, sample XML RPC client implementation that wraps xmlrpclib.ServerProxy and
-    fires locust events on request_success and request_failure, so that all requests
-    gets tracked in locust's statistics.
+    fires locust events on request, so that all requests
+    get tracked in locust's statistics.
     """
 
     _locust_environment = None
@@ -18,20 +18,24 @@ def __getattr__(self, name):
 
         def wrapper(*args, **kwargs):
             start_time = time.time()
+            request_meta = {
+                "request_type": "xmlrpc",
+                "name": name,
+                "response_time": 0,
+                "response_length": 0,
+                "context": {},
+                "exception": None,
+            }
+
             try:
                 result = func(*args, **kwargs)
             except Fault as e:
-                total_time = int((time.time() - start_time) * 1000)
-                self._locust_environment.events.request_failure.fire(
-                    request_type="xmlrpc", name=name, response_time=total_time, exception=e
-                )
-            else:
-                total_time = int((time.time() - start_time) * 1000)
-                self._locust_environment.events.request_success.fire(
-                    request_type="xmlrpc", name=name, response_time=total_time, response_length=0
-                )
-                # In this example, I've hardcoded response_length=0. If we would want the response length to be
-                # reported correctly in the statistics, we would probably need to hook in at a lower level
+                request_meta["exception"] = e
+
+            request_meta["response_time"] = int((time.time() - start_time) * 1000)
+            self._locust_environment.events.request.fire(**request_meta)
+            # In this example, I've hardcoded response_length=0. If we would want the response length to be
+            # reported correctly in the statistics, we would probably need to hook in at a lower level
 
         return wrapper
 

examples/events.py

@@ -47,10 +47,10 @@ def total_content_length():
             return "Total content-length received: %i" % stats["content-length"]
 
 
-@events.request_success.add_listener
-def on_request_success(request_type, name, response_time, response_length):
+@events.request.add_listener
+def on_request(request_type, name, response_time, response_length, exception, context, **kwargs):
     """
-    Event handler that get triggered on every successful request
+    Event handler that get triggered on every request.
     """
     stats["content-length"] += response_length
 

examples/extend_web_ui/extend.py

@@ -123,10 +123,10 @@ def content_length_csv():
         environment.web_ui.app.register_blueprint(extend)
 
 
-@events.request_success.add_listener
-def on_request_success(request_type, name, response_time, response_length):
+@events.request.add_listener
+def on_request(request_type, name, response_time, response_length, exception, context, **kwargs):
     """
-    Event handler that get triggered on every successful request
+    Event handler that get triggered on every request
     """
     stats.setdefault(name, {"content-length": 0})
     stats[name]["content-length"] += response_length

examples/manual_stats_reporting.py

@@ -30,7 +30,7 @@ def _manual_report(name):
     try:
         yield
     except Exception as e:
-        events.request_failure.fire(
+        events.request.fire(
             request_type="manual",
             name=name,
             response_time=(time() - start_time) * 1000,
@@ -39,11 +39,12 @@ def _manual_report(name):
         )
         raise
     else:
-        events.request_success.fire(
+        events.request.fire(
             request_type="manual",
             name=name,
             response_time=(time() - start_time) * 1000,
             response_length=0,
+            exception=None,
         )
 
 

locust/clients.py

@@ -44,12 +44,12 @@ class HttpSession(requests.Session):
                            and then mark it as successful even if the response code was not (i.e 500 or 404).
     """
 
-    def __init__(self, base_url, request_success, request_failure, *args, **kwargs):
+    def __init__(self, base_url, request_event, user, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
         self.base_url = base_url
-        self.request_success = request_success
-        self.request_failure = request_failure
+        self.request_event = request_event
+        self.user = user
 
         # Check for basic authentication
         parsed_url = urlparse(self.base_url)
@@ -66,13 +66,13 @@ def __init__(self, base_url, request_success, request_failure, *args, **kwargs):
             self.auth = HTTPBasicAuth(parsed_url.username, parsed_url.password)
 
     def _build_url(self, path):
-        """ prepend url with hostname unless it's already an absolute URL """
+        """prepend url with hostname unless it's already an absolute URL"""
         if absolute_http_url_regexp.match(path):
             return path
         else:
             return "%s%s" % (self.base_url, path)
 
-    def request(self, method, url, name=None, catch_response=False, **kwargs):
+    def request(self, method, url, name=None, catch_response=False, context={}, **kwargs):
         """
         Constructs and sends a :py:class:`requests.Request`.
         Returns :py:class:`requests.Response` object.
@@ -104,57 +104,46 @@ def request(self, method, url, name=None, catch_response=False, **kwargs):
 
         # prepend url with hostname unless it's already an absolute URL
         url = self._build_url(url)
-
-        # store meta data that is used when reporting the request to locust's statistics
-        request_meta = {}
-
-        # set up pre_request hook for attaching meta data to the request object
-        request_meta["method"] = method
-        request_meta["start_time"] = time.monotonic()
+        start_time = time.monotonic()
 
         response = self._send_request_safe_mode(method, url, **kwargs)
 
-        # record the consumed time
-        request_meta["response_time"] = (time.monotonic() - request_meta["start_time"]) * 1000
+        if self.user:
+            context = {**context, **self.user.context()}
 
-        request_meta["name"] = name or (response.history and response.history[0] or response).request.path_url
+        # store meta data that is used when reporting the request to locust's statistics
+        request_meta = {
+            "request_type": method,
+            "start_time": start_time,
+            "response_time": (time.monotonic() - start_time) * 1000,
+            "name": name or (response.history and response.history[0] or response).request.path_url,
+            "context": context,
+            "exception": None,
+        }
 
         # get the length of the content, but if the argument stream is set to True, we take
         # the size from the content-length header, in order to not trigger fetching of the body
         if kwargs.get("stream", False):
-            request_meta["content_size"] = int(response.headers.get("content-length") or 0)
+            request_meta["response_length"] = int(response.headers.get("content-length") or 0)
         else:
-            request_meta["content_size"] = len(response.content or b"")
+            request_meta["response_length"] = len(response.content or b"")
 
         if catch_response:
-            response.locust_request_meta = request_meta
-            return ResponseContextManager(
-                response, request_success=self.request_success, request_failure=self.request_failure
-            )
+            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:
-                self.request_failure.fire(
-                    request_type=request_meta["method"],
-                    name=request_meta["name"],
-                    response_time=request_meta["response_time"],
-                    response_length=request_meta["content_size"],
-                    exception=e,
-                )
-            else:
-                self.request_success.fire(
-                    request_type=request_meta["method"],
-                    name=request_meta["name"],
-                    response_time=request_meta["response_time"],
-                    response_length=request_meta["content_size"],
-                )
+                request_meta["exception"] = e
+
+            self.request_event.fire(**request_meta)
             if name:
                 response.url = orig_url
             return response
@@ -189,58 +178,45 @@ class ResponseContextManager(LocustResponse):
 
     _manual_result = None
 
-    def __init__(self, response, request_success, request_failure):
+    def __init__(self, response, request_event, request_meta):
         # copy data from response to this object
         self.__dict__ = response.__dict__
-        self._request_success = request_success
-        self._request_failure = request_failure
+        self._request_event = request_event
+        self.request_meta = request_meta
 
     def __enter__(self):
         return self
 
     def __exit__(self, exc, value, traceback):
+        # if the user has already manually marked this response as failure or success
+        # we can ignore the default behaviour of letting the response code determine the outcome
         if self._manual_result is not None:
             if self._manual_result is True:
-                self._report_success()
+                self.request_meta["exception"] = None
             elif isinstance(self._manual_result, Exception):
-                self._report_failure(self._manual_result)
-
-            # if the user has already manually marked this response as failure or success
-            # we can ignore the default behaviour of letting the response code determine the outcome
+                self.request_meta["exception"] = self._manual_result
+            self._report_request()
             return exc is None
 
         if exc:
             if isinstance(value, ResponseError):
-                self._report_failure(value)
+                self.request_meta["exception"] = value
+                self._report_request()
             else:
                 # we want other unknown exceptions to be raised
                 return False
         else:
             try:
                 self.raise_for_status()
             except requests.exceptions.RequestException as e:
-                self._report_failure(e)
-            else:
-                self._report_success()
+                self.request_meta["exception"] = e
+
+            self._report_request()
 
         return True
 
-    def _report_success(self):
-        self._request_success.fire(
-            request_type=self.locust_request_meta["method"],
-            name=self.locust_request_meta["name"],
-            response_time=self.locust_request_meta["response_time"],
-            response_length=self.locust_request_meta["content_size"],
-        )
-
-    def _report_failure(self, exc):
-        self._request_failure.fire(
-            request_type=self.locust_request_meta["method"],
-            name=self.locust_request_meta["name"],
-            response_time=self.locust_request_meta["response_time"],
-            response_length=self.locust_request_meta["content_size"],
-            exception=exc,
-        )
+    def _report_request(self, exc=None):
+        self._request_event.fire(**self.request_meta)
 
     def success(self):
         """