From 62c837751e07c36cf44e39f8243f22f23e2ab617 Mon Sep 17 00:00:00 2001 From: Lars Holmberg Date: Mon, 12 Dec 2022 16:56:44 +0100 Subject: [PATCH 1/3] Add RestUser --- examples/rest_ex.py | 101 +++++++++++++++++++++++++++++++++++++++++ locust/contrib/rest.py | 78 +++++++++++++++++++++++++++++++ tox.ini | 5 +- 3 files changed, 182 insertions(+), 2 deletions(-) create mode 100644 examples/rest_ex.py create mode 100644 locust/contrib/rest.py diff --git a/examples/rest_ex.py b/examples/rest_ex.py new file mode 100644 index 0000000000..957ce14ff5 --- /dev/null +++ b/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) diff --git a/locust/contrib/rest.py b/locust/contrib/rest.py new file mode 100644 index 0000000000..7f4de55d90 --- /dev/null +++ b/locust/contrib/rest.py @@ -0,0 +1,78 @@ +import re +import time +import traceback +from contextlib import contextmanager +from json.decoder import JSONDecodeError +from typing import Generator, Optional, cast + +from locust import FastHttpUser +from locust.clients import ResponseContextManager + + +class RestResponseContextManager(ResponseContextManager): + js: dict # This is technically an Optional, but I dont want to force everyone to check it + error: Exception # This one too + + +class RestUser(FastHttpUser): + """ + A convenience class for testing RESTful JSON endpoints. + 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 (this probably should have been the default behaviour in Locust) + """ + + abstract = True + _callstack_regex = re.compile(r' File "(\/.[^"]*)", line (\d*),(.*)') + + @contextmanager + def rest( + self, method, url, headers: Optional[dict] = None, **kwargs + ) -> Generator[RestResponseContextManager, None, None]: + headers = {"Content-Type": "application/json", "Accept": "application/json"} if headers is None else headers + with self.client.request(method, url, catch_response=True, headers=headers, **kwargs) as r: + resp = cast(RestResponseContextManager, r) + resp.js = None # type: ignore + if resp.text is None: + # round the response time to nearest second to improve error grouping + response_time = round(resp.request_meta["response_time"] / 1000, 1) + resp.failure( + f"response body None, error {resp.error}, response code {resp.status_code}, response time ~{response_time}s" + ) + else: + if resp.text: + try: + resp.js = resp.json() + except JSONDecodeError as e: + resp.failure( + f"Could not parse response as JSON. {resp.text[:250]}, response code {resp.status_code}, error {e}" + ) + try: + yield resp + except AssertionError as e: + if e.args: + if e.args[0].endswith(","): + short_resp = resp.text[:200] if resp.text else resp.text + resp.failure(f"{e.args[0][:-1]}, response was {short_resp}") + else: + resp.failure(e.args[0]) + else: + resp.failure("Assertion failed") + + except Exception as e: + error_lines = [] + for l in traceback.format_exc().split("\n"): + m = self._callstack_regex.match(l) + if m: + filename = re.sub(r"/(home|Users/\w*)/", "~/", m.group(1)) + error_lines.append(filename + ":" + m.group(2) + m.group(3)) + short_resp = resp.text[:200] if resp.text else resp.text + resp.failure(f"{e.__class__.__name__}: {e} at {', '.join(error_lines)}. Response was {short_resp}") + + # some web api:s use a timestamp as part of their url (to break thru caches). this is a convenience method for that. + @contextmanager + def rest_(self, method, url, name=None, **kwargs) -> Generator[RestResponseContextManager, None, None]: + with self.rest(method, f"{url}&_={int(time.time()*1000)}", name=name, **kwargs) as resp: + yield resp diff --git a/tox.ini b/tox.ini index 2f0f94ed4e..290f01ce93 100644 --- a/tox.ini +++ b/tox.ini @@ -28,8 +28,9 @@ allowlist_externals = commands = python3 -m pip install . python3 -m coverage run -m unittest discover [] - bash -ec 'PYTHONUNBUFFERED=1 timeout 2s python3 examples/debugging.py >out.txt 2>err.txt || true' - grep -qm 1 '/hello' out.txt + bash -ec "PYTHONUNBUFFERED=1 timeout 10s python3 examples/rest_ex.py >out.txt 2>err.txt || true" + grep -qm 1 'my custom error message with response text, response was {"args"' out.txt + grep -qm 1 'ZeroDivisionError: division by zero at.*Response was {"ar' out.txt bash -ec '! grep . err.txt' # should be empty bash -ec 'PYTHONUNBUFFERED=1 python3 examples/debugging_advanced.py | grep done' From 213959a2ce73a69fdb8f9dbdd4e8ee73ffd66417 Mon Sep 17 00:00:00 2001 From: Lars Holmberg Date: Mon, 12 Dec 2022 20:29:03 +0100 Subject: [PATCH 2/3] Add RestUser to documentation. Rework testing-other-systems docs. Rework gRPC example code (mainly moving the ClientInterceptor and abstract base User into its own file) --- docs/index.rst | 1 - docs/testing-other-systems.rst | 74 +++++++++++++----- docs/testing-requests-based SDK's.rst | 18 ----- examples/grpc/grpc_user.py | 61 +++++++++++++++ examples/grpc/locustfile.py | 76 ++----------------- .../session_patch_locustfile.py | 4 +- 6 files changed, 123 insertions(+), 111 deletions(-) delete mode 100644 docs/testing-requests-based SDK's.rst create mode 100644 examples/grpc/grpc_user.py diff --git a/docs/index.rst b/docs/index.rst index 40e30c7e8f..2a504001c9 100644 --- a/docs/index.rst +++ b/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 diff --git a/docs/testing-other-systems.rst b/docs/testing-other-systems.rst index c93dde5112..c5fe23876a 100644 --- a/docs/testing-other-systems.rst +++ b/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 ` event after each call has completed, to let Locust know what happened. @@ -14,10 +14,10 @@ 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 `_. 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 @@ -25,27 +25,63 @@ We can build a generic XML-RPC client, by wrapping :py:class:`xmlrpc.client.Serv .. 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 `_ User. +We can also build a `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 `_: - grpc_gevent.init_gevent() +.. literalinclude:: ../examples/grpc/grpc_user.py -Dummy server to test: +And an example test 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 `_ 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 `_. In that case, just pass it your ``HttpUser``'s :py:attr:`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 `_ 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 `. 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 `_. That also shows how you can subclass :py:class:`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 ` library. + For more examples of user types, see `locust-plugins `_ (it has users for WebSocket/SocketIO, Kafka, Selenium/WebDriver and more). -For more examples of user types, see `locust-plugins `_ (it has users for WebSocket/SocketIO, Kafka, Selenium/WebDriver and more). \ No newline at end of file diff --git a/docs/testing-requests-based SDK's.rst b/docs/testing-requests-based SDK's.rst deleted file mode 100644 index b52ac579f4..0000000000 --- a/docs/testing-requests-based SDK's.rst +++ /dev/null @@ -1,18 +0,0 @@ -.. _testing-request-sdks: - -============================= -Testing Requests based SDKs -============================= - -If a prebuilt SDK is available for your target system, Locust has a supported pattern for integrating -its usage into your load testing efforts. - -The only prerequisite to achieve this is that the SDK needs to have an accessible ``request.Sessions`` -class. - -The following example shows the locust client overwriting the internal ``_session`` object of ``Archivist`` SDK -during startup. - -.. literalinclude:: ../examples/sdk_session_patching/session_patch_locustfile.py - - diff --git a/examples/grpc/grpc_user.py b/examples/grpc/grpc_user.py new file mode 100644 index 0000000000..08ba004eff --- /dev/null +++ b/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) diff --git a/examples/grpc/locustfile.py b/examples/grpc/locustfile.py index c3cbee0d69..a10a514d69 100644 --- a/examples/grpc/locustfile.py +++ b/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) diff --git a/examples/sdk_session_patching/session_patch_locustfile.py b/examples/sdk_session_patching/session_patch_locustfile.py index e971219ae9..90471e7e7f 100644 --- a/examples/sdk_session_patching/session_patch_locustfile.py +++ b/examples/sdk_session_patching/session_patch_locustfile.py @@ -1,6 +1,6 @@ import locust from locust.user import task -from archivist.archivist import Archivist # Example SDK under test +from archivist.archivist import Archivist # Example library under test class ArchivistUser(locust.HttpUser): @@ -10,7 +10,7 @@ def on_start(self): with open("auth.text") as f: AUTH_TOKEN = f.read() - # Start an instance of of the SDK + # Start an instance of of the library-provided client self.arch: Archivist = Archivist(url=self.host, auth=AUTH_TOKEN) # overwrite the internal _session attribute with the locust session self.arch._session = self.client From 5a45e28a1b6b4508c408847d23d7baa69f59b31d Mon Sep 17 00:00:00 2001 From: Lars Holmberg Date: Mon, 12 Dec 2022 20:29:19 +0100 Subject: [PATCH 3/3] docs --- docs/testing-other-systems.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/testing-other-systems.rst b/docs/testing-other-systems.rst index c5fe23876a..98a7be53a2 100644 --- a/docs/testing-other-systems.rst +++ b/docs/testing-other-systems.rst @@ -38,7 +38,7 @@ The generic gRPC User base class sends events to Locust using an `interceptor