Merge pull request #2273 from locustio/add-RestUser

Add RestUser

docs/index.rst

@@ -54,7 +54,6 @@ Other functionalities
     custom-load-shape
     retrieving-stats
     testing-other-systems
-    testing-requests-based SDK's
     increase-performance
     extending-locust
     logging

docs/testing-other-systems.rst

@@ -1,8 +1,8 @@
 .. _testing-other-systems:
 
-========================
-Testing non-HTTP systems
-========================
+===============================
+Testing other systems/protocols
+===============================
 
 Locust only comes with built-in support for HTTP/HTTPS but it can be extended to test almost any system. This is normally done by wrapping the protocol library and triggering a :py:attr:`request <locust.event.Events.request>` event after each call has completed, to let Locust know what happened.
 
@@ -14,38 +14,74 @@ Locust only comes with built-in support for HTTP/HTTPS but it can be extended to
 
     Some C libraries 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/>`_. If you are willing to get your hands dirty, you may also be able to patch a library yourself, but that is beyond the scope of this documentation.
 
-Example: writing an XML-RPC User/client
-=======================================
+Example: XML-RPC
+================
 
-Lets assume we had an XML-RPC server that we wanted to load test.
+Lets assume we have an XML-RPC server that we want to load test.
 
 .. literalinclude:: ../examples/custom_xmlrpc_client/server.py
 
 We can build a generic XML-RPC client, by wrapping :py:class:`xmlrpc.client.ServerProxy`.
 
 .. literalinclude:: ../examples/custom_xmlrpc_client/xmlrpc_locustfile.py
 
-Example: writing a gRPC User/client
-=======================================
+Example: gRPC
+=============
 
-If you have understood the XML-RPC example, you can easily build a `gRPC <https://github.com/grpc/grpc>`_ User.
+We can also build a `gRPC <https://github.com/grpc/grpc>`_ User.
 
-The only significant difference is that you need to make gRPC gevent-compatible, by executing this code before opening the channel:
+Lets assume we have a gRPC server that we want to load test:
 
-.. code-block:: python
+.. literalinclude:: ../examples/grpc/hello_server.py
 
-    import grpc.experimental.gevent as grpc_gevent
+The generic gRPC User base class sends events to Locust using an `interceptor <https://pypi.org/project/grpc-interceptor/>`_:
 
-    grpc_gevent.init_gevent()
+.. literalinclude:: ../examples/grpc/grpc_user.py
 
-Dummy server to test:
+And a locustfile using the above would look like this:
 
-.. literalinclude:: ../examples/grpc/hello_server.py
+.. literalinclude:: ../examples/grpc/locustfile.py
 
-gRPC client, base GrpcUser, interceptor for sending events to locust and example usage:
+.. _testing-request-sdks:
 
-.. literalinclude:: ../examples/grpc/locustfile.py
+requests-based libraries/SDKs
+=============================
+
+If you want to use a library that uses a `requests.Session <https://requests.readthedocs.io/en/latest/api/#requests.Session>`_ object under the hood you will most likely be able to skip all the above complexity.
+
+Some libraries allow you to pass a Session explicitly, like for example the SOAP client provided by `Zeep <https://docs.python-zeep.org/en/master/transport.html#tls-verification>`_. In that case, just pass it your ``HttpUser``'s :py:attr:`client <locust.HttpUser.client>`, and any requests made using the library will be logged in Locust.
+
+Even if your library doesn't expose that in its interface, you may be able to get it working by overwriting some internally used Session. Here's an example of how to do that for the `Archivist <https://github.com/jitsuin-inc/archivist-python>`_ client.
+
+.. literalinclude:: ../examples/sdk_session_patching/session_patch_locustfile.py
+
+
+Example: REST
+=============
+
+While the base HttpUser/FastHttpUser is capable of testing RESTful endpoints, it can be simplified by using a specialized subclass :py:class:`RestUser <locust.contrib.rest.RestUser>`. It extends FastHttpUser by adding the ``rest``-method, a wrapper around ``self.client.request()`` that:
+
+* automatically passes ``catch_response=True``
+* automatically sets ``Content-Type`` and ``Accept`` headers to ``application/json`` (unless you have provided your own headers)
+* automatically checks that the response is valid json, parses it into an dict and saves it in a field called ``js`` in the response object.
+* catches any exceptions thrown in your ``with``-block and fails the sample (instead of crashing the task)
+
+.. code-block:: python
+
+    from locust.contrib.rest import RestUser
+    from locust import task
+
+    class MyUser(RestUser):
+        @task
+        def t(self):
+            with self.rest("POST", "/", json={"foo": 1}) as resp:
+                if resp.js and resp.js["bar"] != 1:
+                    resp.failure(f"Unexpected value of foo in response {resp.text}")
+
+For a complete example, see `resp_ex.py <https://github.com/locustio/locust/tree/master/examples/resp_ex.py>`_. That also shows how you can subclass :py:class:`RestUser <locust.contrib.rest.RestUser>` to provide behaviours specific to your API, like like always sending common headers or always applying some validation to the response.
+
+
+.. note::
 
-As base class for interceptor is used `grpc-interceptor <https://pypi.org/project/grpc-interceptor/>` library.
+    For more examples of user types, see `locust-plugins <https://github.com/SvenskaSpel/locust-plugins#users>`_ (it has users for WebSocket/SocketIO, Kafka, Selenium/WebDriver and more).
 
-For more examples of user types, see `locust-plugins <https://github.com/SvenskaSpel/locust-plugins#users>`_ (it has users for WebSocket/SocketIO, Kafka, Selenium/WebDriver and more).

examples/grpc/grpc_user.py

@@ -0,0 +1,61 @@
+import time
+from typing import Any, Callable
+import grpc
+import grpc.experimental.gevent as grpc_gevent
+from grpc_interceptor import ClientInterceptor
+from locust import User
+from locust.exception import LocustError
+
+# patch grpc so that it uses gevent instead of asyncio
+grpc_gevent.init_gevent()
+
+
+class LocustInterceptor(ClientInterceptor):
+    def __init__(self, environment, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.env = environment
+
+    def intercept(
+        self,
+        method: Callable,
+        request_or_iterator: Any,
+        call_details: grpc.ClientCallDetails,
+    ):
+        response = None
+        exception = None
+        start_perf_counter = time.perf_counter()
+        response_length = 0
+        try:
+            response = method(request_or_iterator, call_details)
+            response_length = response.result().ByteSize()
+        except grpc.RpcError as e:
+            exception = e
+
+        self.env.events.request.fire(
+            request_type="grpc",
+            name=call_details.method,
+            response_time=(time.perf_counter() - start_perf_counter) * 1000,
+            response_length=response_length,
+            response=response,
+            context=None,
+            exception=exception,
+        )
+        return response
+
+
+class GrpcUser(User):
+    abstract = True
+    stub_class = None
+
+    def __init__(self, environment):
+        super().__init__(environment)
+        for attr_value, attr_name in ((self.host, "host"), (self.stub_class, "stub_class")):
+            if attr_value is None:
+                raise LocustError(f"You must specify the {attr_name}.")
+
+        self._channel = grpc.insecure_channel(self.host)
+        interceptor = LocustInterceptor(environment=environment)
+        self._channel = grpc.intercept_channel(self._channel, interceptor)
+
+        self.stub = self.stub_class(self._channel)

examples/grpc/locustfile.py

@@ -1,87 +1,21 @@
-# make sure you use grpc version 1.39.0 or later,
-# because of https://github.com/grpc/grpc/issues/15880 that affected earlier versions
-from typing import Callable, Any
-import time
-
-import grpc
-import grpc.experimental.gevent as grpc_gevent
 import gevent
-from locust import events, User, task
-from locust.exception import LocustError
-from grpc_interceptor import ClientInterceptor
-
-import hello_pb2_grpc
+import grpc_user
 import hello_pb2
-
+import hello_pb2_grpc
 from hello_server import start_server
-
-# patch grpc so that it uses gevent instead of asyncio
-grpc_gevent.init_gevent()
+from locust import events, task
 
 
+# Start the dummy server. This is not something you would do in a real test.
 @events.init.add_listener
 def run_grpc_server(environment, **_kwargs):
-    # Start the dummy server. This is not something you would do in a real test.
     gevent.spawn(start_server)
 
 
-class LocustInterceptor(ClientInterceptor):
-    def __init__(self, environment, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-        self.env = environment
-
-    def intercept(
-        self,
-        method: Callable,
-        request_or_iterator: Any,
-        call_details: grpc.ClientCallDetails,
-    ):
-        response = None
-        exception = None
-        start_perf_counter = time.perf_counter()
-        response_length = 0
-        try:
-            response = method(request_or_iterator, call_details)
-            response_length = response.result().ByteSize()
-        except grpc.RpcError as e:
-            exception = e
-
-        self.env.events.request.fire(
-            request_type="grpc",
-            name=call_details.method,
-            response_time=(time.perf_counter() - start_perf_counter) * 1000,
-            response_length=response_length,
-            response=response,
-            context=None,
-            exception=exception,
-        )
-        return response
-
-
-class GrpcUser(User):
-    abstract = True
-
-    stub_class = None
-
-    def __init__(self, environment):
-        super().__init__(environment)
-        for attr_value, attr_name in ((self.host, "host"), (self.stub_class, "stub_class")):
-            if attr_value is None:
-                raise LocustError(f"You must specify the {attr_name}.")
-
-        self._channel = grpc.insecure_channel(self.host)
-        interceptor = LocustInterceptor(environment=environment)
-        self._channel = grpc.intercept_channel(self._channel, interceptor)
-
-        self.stub = self.stub_class(self._channel)
-
-
-class HelloGrpcUser(GrpcUser):
+class HelloGrpcUser(grpc_user.GrpcUser):
     host = "localhost:50051"
     stub_class = hello_pb2_grpc.HelloServiceStub
 
     @task
     def sayHello(self):
         self.stub.SayHello(hello_pb2.HelloRequest(name="Test"))
-        time.sleep(1)

examples/rest_ex.py

@@ -0,0 +1,101 @@
+from contextlib import contextmanager
+from locust import task, run_single_user
+from locust.contrib.fasthttp import ResponseContextManager
+from locust.user.wait_time import constant
+from locust.contrib.rest import RestUser
+
+
+class MyUser(RestUser):
+    host = "https://postman-echo.com"
+    wait_time = constant(180)  # be nice to postman-echo.com, and dont run this at scale.
+
+    @task
+    def t(self):
+        # should work
+        with self.rest("GET", "/get", json={"foo": 1}) as resp:
+            if resp.js["args"]["foo"] != 1:
+                resp.failure(f"Unexpected value of foo in response {resp.text}")
+
+        # should work
+        with self.rest("POST", "/post", json={"foo": 1}) as resp:
+            if resp.js["data"]["foo"] != 1:
+                resp.failure(f"Unexpected value of foo in response {resp.text}")
+            # assertions are a nice short way of expressiont your expectations about the response. The AssertionError thrown will be caught
+            # and fail the request, including the message and the payload in the failure content
+            assert resp.js["data"]["foo"] == 1, "Unexpected value of foo in response"
+
+        # assertions are a nice short way to validate the response. The AssertionError they raise
+        # will be caught by rest() and mark the request as failed
+
+        with self.rest("POST", "/post", json={"foo": 1}) as resp:
+            # mark the request as failed with the message "Assertion failed"
+            assert resp.js["data"]["foo"] == 2
+
+        with self.rest("POST", "/post", json={"foo": 1}) as resp:
+            # custom failure message
+            assert resp.js["data"]["foo"] == 2, "my custom error message"
+
+        with self.rest("POST", "/post", json={"foo": 1}) as resp:
+            # use a trailing comma to append the response text to the custom message
+            assert resp.js["data"]["foo"] == 2, "my custom error message with response text,"
+
+        # this only works in python 3.8 and up, so it is commented out:
+        # if sys.version_info >= (3, 8):
+        #     with self.rest("", "/post", json={"foo": 1}) as resp:
+        #         # assign and assert in one line
+        #         assert (foo := resp.js["foo"])
+        #         print(f"the number {foo} is awesome")
+
+        # rest() catches most exceptions, so any programming mistakes you make automatically marks the request as a failure
+        # and stores the callstack in the failure message
+        with self.rest("POST", "/post", json={"foo": 1}) as resp:
+            1 / 0  # pylint: disable=pointless-statement
+
+        # response isnt even json, but RestUser will already have been marked it as a failure, so we dont have to do it again
+        with self.rest("GET", "/") as resp:
+            pass
+
+        with self.rest("GET", "/") as resp:
+            # If resp.js is None (which it will be when there is a connection failure, a non-json responses etc),
+            # reading from resp.js will raise a TypeError (instead of an AssertionError), so lets avoid that:
+            if resp.js:
+                assert resp.js["foo"] == 2
+            # or, as a mildly confusing oneliner:
+            assert not resp.js or resp.js["foo"] == 2
+
+        # 404
+        with self.rest("GET", "http://example.com/") as resp:
+            pass
+
+        # connection closed
+        with self.rest("POST", "http://example.com:42/", json={"foo": 1}) as resp:
+            pass
+
+
+# An example of how you might write a common base class for an API that always requires
+# certain headers, or where you always want to check the response in a certain way
+class RestUserThatLooksAtErrors(RestUser):
+    abstract = True
+
+    @contextmanager
+    def rest(self, method, url, **kwargs) -> ResponseContextManager:
+        extra_headers = {"my_header": "my_value"}
+        with super().rest(method, url, headers=extra_headers, **kwargs) as resp:
+            resp: ResponseContextManager
+            if resp.js and "error" in resp.js and resp.js["error"] is not None:
+                resp.failure(resp.js["error"])
+            yield resp
+
+
+class MyOtherRestUser(RestUserThatLooksAtErrors):
+    host = "https://postman-echo.com"
+    wait_time = constant(180)  # be nice to postman-echo.com, and dont run this at scale.
+
+    @task
+    def t(self):
+        with self.rest("GET", "/") as _resp:
+            pass
+
+
+if __name__ == "__main__":
+    run_single_user(MyUser)