Merge pull request #1752 from locustio/add-response-object-to-request…

…-event

Add response object to request event

docs/changelog.rst

@@ -7,8 +7,8 @@ For full details of the Locust changelog, please see https://github.com/locustio
 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
+* Unify request_success/request_failure into a single event called request (the old ones are deprecated but still work) https://github.com/locustio/locust/issues/1724
+* Add the response object and context as parameters to the request event. context is used to forward information to the request event handler (can be used for things like username, tags etc)
 
 1.4.4
 =====

docs/extending-locust.rst

@@ -16,19 +16,23 @@ Here's an example on how to set up an event listener::
     from locust import events
     
     @events.request.add_listener
-    def my_request_handler(request_type, name, response_time, response_length, context, exception, **kw):
+    def my_request_handler(request_type, name, response_time, response_length, response,
+                           context, exception, **kw):
         if exception:
             print(f"Request to {name} failed with exception {exception}")
         else:
             print(f"Successfully made a request to: {name})
 
-
 .. note::
 
     It's highly recommended that you add a wildcard keyword argument in your listeners
     (the \**kw in the code above), to prevent your code from breaking if new arguments are
     added in a future version.
 
+    Note that it is entirely possible to implement a client that does not support all parameters 
+    (some non-HTTP protocols might not have a concept of `response_length` or `response` object).
+
+.. _request_context:
 
 Request context
 ==================
@@ -60,7 +64,7 @@ Context from User class::
             self.client.post("/login", json={"username": self.username})
 
         @events.request.add_listener
-        def on_request(self, context, **kwargs):
+        def on_request(context, **kwargs):
             print(context["username"])
 
 

docs/testing-other-systems.rst

@@ -4,38 +4,23 @@
 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 <locust.event.Events.request>`
+Locust was built with HTTP as its main use case but it can be extended to load test almost any system. You do this by writing a custom client that triggers :py:attr:`request <locust.event.Events.request>`
 
 .. note::
 
-    Any protocol libraries that you use must be gevent-friendly (use the Python ``socket`` module or some other standard library function like ``subprocess``), or your calls will block the whole Locust process.
+    Any protocol libraries that you use must be gevent-friendly (use the Python ``socket`` module or some other standard library function like ``subprocess``), or your calls are likely to block the whole Locust/Python process.
 
-    Some C libraries cannot be monkey patched by gevent, but allow for other workarounds. For example, if you want to use psycopg2 to performance test PostgreSQL, can use `psycogreen <https://github.com/psycopg/psycogreen/>`_. 
+    Some C libraries cannot be monkey patched by gevent, but allow for other workarounds. For example, if you want to use psycopg2 to performance test PostgreSQL, you can use `psycogreen <https://github.com/psycopg/psycogreen/>`_. 
 
-Sample XML-RPC User client
-============================
+Example: writing an XML-RPC User/client
+=======================================
 
-Here is an example of a User class, **XmlRpcUser**, which provides an XML-RPC client,
-**XmlRpcUser**, and tracks all requests made:
+Lets assume we had an XML-RPC server that we wanted to load test
 
-.. literalinclude:: ../examples/custom_xmlrpc_client/xmlrpc_locustfile.py
-
-If you've written Locust tests before, you'll recognize the class called ``ApiUser`` which is a normal 
-User class that has a couple of tasks declared. However, the ``ApiUser`` inherits from
-``XmlRpcUser`` that you can see right above ``ApiUser``. The ``XmlRpcUser`` is marked as abstract
-using ``abstract = True`` which means that Locust will not try to create simulated users from that class
-(only of classes that extend it). ``XmlRpcUser`` provides an instance of XmlRpcClient under the
-``client`` attribute. 
-
-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`
-event, which will record all calls in Locust's statistics.
+.. literalinclude:: ../examples/custom_xmlrpc_client/server.py
 
-Here's an implementation of an XML-RPC server that would work as a server for the code above:
+We can build a generic XML-RPC client, by wrapping :py:class:`xmlrpc.client.ServerProxy`
 
-.. literalinclude:: ../examples/custom_xmlrpc_client/server.py
+.. literalinclude:: ../examples/custom_xmlrpc_client/xmlrpc_locustfile.py
 
 For more examples, see `locust-plugins <https://github.com/SvenskaSpel/locust-plugins#users>`_

examples/custom_xmlrpc_client/xmlrpc_locustfile.py

@@ -1,67 +1,65 @@
 import time
 from xmlrpc.client import ServerProxy, Fault
 
-from locust import User, task, between
+from locust import User, task
 
 
 class XmlRpcClient(ServerProxy):
     """
-    Simple, sample XML RPC client implementation that wraps xmlrpclib.ServerProxy and
-    fires locust events on request, so that all requests
-    get tracked in locust's statistics.
+    XmlRpcClient is a wrapper around the standard library's ServerProxy.
+    It proxies any function calls and fires the *request* event when they finish,
+    so that the calls get recorded in Locust.
     """
 
-    _locust_environment = None
+    def __init__(self, host, request_event):
+        super().__init__(host)
+        self._request_event = request_event
 
     def __getattr__(self, name):
         func = ServerProxy.__getattr__(self, name)
 
         def wrapper(*args, **kwargs):
-            start_time = time.time()
+            start_time = time.monotonic()
             request_meta = {
                 "request_type": "xmlrpc",
                 "name": name,
-                "response_time": 0,
-                "response_length": 0,
-                "context": {},
+                "response_length": 0,  # calculating this for an xmlrpc.client response would be too hard
+                "response": None,
+                "context": {},  # see HttpUser if you actually want to implement contexts
                 "exception": None,
             }
-
             try:
-                result = func(*args, **kwargs)
+                request_meta["response"] = func(*args, **kwargs)
             except Fault as e:
                 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
+            request_meta["response_time"] = (time.monotonic() - start_time) * 1000
+            self._request_event.fire(**request_meta)  # This is what makes the request actually get logged in Locust
+            return request_meta["response"]
 
         return wrapper
 
 
 class XmlRpcUser(User):
     """
-    This is the abstract User class which should be subclassed. It provides an XML-RPC client
-    that can be used to make XML-RPC requests that will be tracked in Locust's statistics.
+    A minimal Locust user class that provides an XmlRpcClient to its subclasses
     """
 
-    abstract = True
+    abstract = True  # dont instantiate this as an actual user when running Locust
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.client = XmlRpcClient(self.host)
-        self.client._locust_environment = self.environment
+    def __init__(self, environment):
+        super().__init__(environment)
+        self.client = XmlRpcClient(self.host, request_event=environment.events.request)
 
 
-class ApiUser(XmlRpcUser):
+# The real user class that will be instantiated and run by Locust
+# This is the only thing that is actually specific to the service that we are testing.
+class MyUser(XmlRpcUser):
     host = "http://127.0.0.1:8877/"
-    wait_time = between(0.1, 1)
 
-    @task(10)
+    @task
     def get_time(self):
         self.client.get_time()
 
-    @task(5)
+    @task
     def get_random_number(self):
         self.client.get_random_number(0, 100)

locust/clients.py

@@ -109,7 +109,7 @@ def request(self, method, url, name=None, catch_response=False, context={}, **kw
         response = self._send_request_safe_mode(method, url, **kwargs)
 
         if self.user:
-            context = {**context, **self.user.context()}
+            context = {**self.user.context(), **context}
 
         # store meta data that is used when reporting the request to locust's statistics
         request_meta = {
@@ -118,6 +118,7 @@ def request(self, method, url, name=None, catch_response=False, context={}, **kw
             "response_time": (time.monotonic() - start_time) * 1000,
             "name": name or (response.history and response.history[0] or response).request.path_url,
             "context": context,
+            "response": response,
             "exception": None,
         }
 

locust/contrib/fasthttp.py

@@ -164,7 +164,7 @@ def request(
         start_time = default_timer()
 
         if self.user:
-            context = {**context, **self.user.context()}
+            context = {**self.user.context(), **context}
 
         # store meta data that is used when reporting the request to locust's statistics
         request_meta = {
@@ -196,6 +196,7 @@ def request(
 
         # send request, and catch any exceptions
         response = self._send_request_safe_mode(method, url, payload=data, headers=headers, **kwargs)
+        request_meta["response"] = response
 
         if not allow_redirects:
             self.client.redirect_resonse_codes = old_redirect_response_codes

locust/event.py

@@ -68,8 +68,9 @@ class Events:
     :param name: Path to the URL that was called (or override name if it was used in the call to the client)
     :param response_time: Time in milliseconds until exception was thrown
     :param response_length: Content-length of the response
-    :param exception: Exception instance that was thrown. None if no exception
-    :param context: Dict with context values specified when performing request
+    :param response: Response object (e.g. a :py:class:`requests.Response`)
+    :param context: :ref:`User/request context <request_context>`
+    :param exception: Exception instance that was thrown. None if request was successful.
     """
 
     request_success: DeprecatedEventHook

locust/test/test_fasthttp.py

@@ -255,7 +255,7 @@ class MyUser(FastHttpUser):
             host = "http://127.0.0.1:%i" % self.port
 
             def context(self):
-                return {"user": self}
+                return {"user": self.username}
 
         kwargs = {}
 
@@ -264,8 +264,12 @@ def on_request(**kw):
 
         self.environment.events.request.add_listener(on_request)
         user = MyUser(self.environment)
+        user.username = "foo"
         user.client.request("get", "/request_method")
-        self.assertDictEqual({"user": user}, kwargs["context"])
+        self.assertDictEqual({"user": "foo"}, kwargs["context"])
+        self.assertEqual("GET", kwargs["response"].text)
+        user.client.request("get", "/request_method", context={"user": "bar"})
+        self.assertDictEqual({"user": "bar"}, kwargs["context"])
 
     def test_get_request(self):
         self.response = ""

locust/test/test_http.py

@@ -132,6 +132,19 @@ def on_request(exception, **kw):
         s.request("get", "/request_method", context={"foo": "bar"})
         self.assertDictEqual({"foo": "bar"}, kwargs["context"])
 
+    def test_response_parameter(self):
+        s = self.get_client()
+        kwargs = {}
+
+        def on_request(**kw):
+            kwargs.update(kw)
+
+        self.environment.events.request.add_listener(on_request)
+        s.request("get", "/request_method")
+        self.assertEqual("GET", kwargs["response"].text)
+        s.request("get", "/wrong_url")
+        self.assertEqual("Not Found", kwargs["response"].text)
+
     def test_deprecated_request_events(self):
         s = self.get_client()
         status = {"success_amount": 0, "failure_amount": 0}
@@ -243,10 +256,10 @@ def test_catch_response_default_fail(self):
 
     def test_user_context(self):
         class TestUser(HttpUser):
-            host = "http://localhost"
+            host = f"http://127.0.0.1:{self.port}"
 
             def context(self):
-                return {"user": self}
+                return {"user": self.username}
 
         kwargs = {}
 
@@ -256,5 +269,9 @@ def on_request(**kw):
         self.environment.events.request.add_listener(on_request)
 
         user = TestUser(self.environment)
+        user.username = "foo"
         user.client.request("get", "/request_method")
-        self.assertDictEqual({"user": user}, kwargs["context"])
+        self.assertDictEqual({"user": "foo"}, kwargs["context"])
+        self.assertEqual("GET", kwargs["response"].text)
+        user.client.request("get", "/request_method", context={"user": "bar"})  # override User context
+        self.assertDictEqual({"user": "bar"}, kwargs["context"])

locust/user/users.py

@@ -187,10 +187,7 @@ def stop(self, force=False):
 
     def context(self) -> Dict:
         """
-        Returns user specific context. Override this method to customize data to be forwarded in request event.
-
-        :return: Context data
-        :rtype: Dict
+        Adds the returned value (a dict) to the context for :ref:`request event <request_context>`
         """
         return {}