Merge pull request #2274 from locustio/integrate-RestUser-into-FastHt…

…tpUser-instead

Move the rest method into FastHttpUser instead of RestUser.

docs/api.rst

@@ -19,14 +19,9 @@ FastHttpUser class
 ==================
 
 .. autoclass:: locust.contrib.fasthttp.FastHttpUser
+    :members: wait_time, tasks, client, abstract, rest
     :noindex:
 
-RestUser class
-================
-
-.. autoclass:: locust.contrib.rest.RestUser
-    :members: rest
-
 TaskSet class
 =============
 

docs/increase-performance.rst

@@ -51,6 +51,40 @@ A single FastHttpUser/geventhttpclient session can run concurrent requests, you
 
     FastHttpUser/geventhttpclient is very similar to HttpUser/python-requests, but sometimes there are subtle differences. This is particularly true if you work with the client library's internals, e.g. when manually managing cookies.
 
+.. _rest:
+
+REST
+====
+
+FastHttpUser also provides a ``rest`` method for testing REST/JSON HTTP interfaces:. It is a wrapper for ``self.client.request`` that:
+
+* Parses the JSON response to a dict called ``js`` in the response object. Marks the request as failed if the response was not valid JSON.
+* Defaults ``Content-Type`` and ``Accept`` headers to ``application/json``
+* Sets ``catch_response=True`` (so always use a :ref:`with-block <catch-response>`)
+* Catches any unhandled exceptions thrown inside your with-block, marking the sample as failed (instead of exiting the task immediately without even firing the request event)
+
+.. code-block:: python
+
+    from locust import task, FastHttpUser
+
+    class MyUser(FastHttpUser):
+        @task
+        def t(self):
+            with self.rest("POST", "/", json={"foo": 1}) as resp:
+                if resp.js is None:
+                    pass # no need to do anything, already marked as failed
+                elif "bar" not in resp.js:
+                    resp.failure(f"'bar' missing from response {resp.text}")
+                elif resp.js["bar"] != 42:
+                    resp.failure(f"'bar' had an unexpected value: {resp.js['bar']}")
+
+For a complete example, see `rest.py <https://github.com/locustio/locust/blob/master/examples/rest.py>`_. That also shows how you can use inheritance to provide behaviours specific to your REST API that are common to multiple requests/testplans.
+
+.. note::
+
+    This feature is new and details of its interface/implementation may change in new versions of Locust.
+
+
 API
 ===
 

docs/testing-other-systems.rst

@@ -53,38 +53,10 @@ Even if your library doesn't expose that in its interface, you may be able to ge
 
 .. literalinclude:: ../examples/sdk_session_patching/session_patch_locustfile.py
 
-
 REST
 ====
 
-While you can subclass :ref:`HttpUser <quickstart>`/:ref:`FastHttpUser <increase-performance>` to test RESTful HTTP endpoints, you can avoid having to reinvent the wheel by using :py:class:`RestUser <locust.contrib.rest.RestUser>`. It extends FastHttpUser, adding the ``rest`` method, a wrapper around :py:class:`self.client.request <locust.contrib.fasthttp.FastHttpUser.client>` that:
-
-* Parses the JSON response to a dict called ``js`` in the response object. Marks the request as failed if the response was not valid JSON.
-* Defaults ``Content-Type`` and ``Accept`` headers to ``application/json``
-* Sets ``catch_response=True`` (so use a :ref:`with-block <catch-response>`)
-* Catches any unhandled exceptions thrown inside your with-block, marking the sample as failed (instead of exiting the task immediately without even firing the request event)
-
-.. code-block:: python
-
-    from locust import task, RestUser
-
-    class MyUser(RestUser):
-        @task
-        def t(self):
-            with self.rest("POST", "/", json={"foo": 1}) as resp:
-                if resp.js is None:
-                    pass # no need to do anything, already marked as failed
-                elif "bar" not in resp.js:
-                    resp.failure(f"'bar' missing from response {resp.text}")
-                elif resp.js["bar"] != 42:
-                    resp.failure(f"'bar' had an unexpected value: {resp.js['bar']}")
-
-For a complete example, see `resp_ex.py <https://github.com/locustio/locust/blob/master/examples/rest_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::
-
-    RestUser is new and details of its interface/implementation may change in new versions of Locust.
-
+See :ref:`FastHttpUser <rest>`
 
 Other examples
 ==============

examples/rest_ex.py → examples/rest.py

@@ -1,11 +1,11 @@
 from contextlib import contextmanager
-from locust import task, run_single_user, RestUser
-from locust.contrib.rest import RestResponseContextManager
+from locust import task, run_single_user, FastHttpUser
+from locust.contrib.fasthttp import RestResponseContextManager
 from locust.user.wait_time import constant
 from typing import Generator
 
 
-class MyUser(RestUser):
+class MyUser(FastHttpUser):
     host = "https://postman-echo.com"
     wait_time = constant(180)  # be nice to postman-echo.com, and dont run this at scale.
 
@@ -74,7 +74,7 @@ def t(self):
 
 # 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):
+class RestUserThatLooksAtErrors(FastHttpUser):
     abstract = True
 
     @contextmanager

locust/__init__.py

@@ -16,7 +16,6 @@
 from .user.task import task, tag, TaskSet
 from .user.users import HttpUser, User
 from .contrib.fasthttp import FastHttpUser
-from .contrib.rest import RestUser
 from .user.wait_time import between, constant, constant_pacing, constant_throughput
 from .shape import LoadTestShape
 from .debug import run_single_user
@@ -33,7 +32,6 @@
     "TaskSet",
     "HttpUser",
     "FastHttpUser",
-    "RestUser",
     "User",
     "between",
     "constant",

locust/contrib/fasthttp.py

@@ -4,10 +4,13 @@
 import json
 import json as unshadowed_json  # some methods take a named parameter called json
 from base64 import b64encode
+from contextlib import contextmanager
+from json.decoder import JSONDecodeError
 from urllib.parse import urlparse, urlunparse
 from ssl import SSLError
 import time
-from typing import Callable, Optional, Tuple, Dict, Any
+import traceback
+from typing import Callable, Optional, Tuple, Dict, Any, Generator, cast
 
 from http.cookiejar import CookieJar
 
@@ -323,6 +326,8 @@ class FastHttpUser(User):
     abstract = True
     """Dont register this as a User class that can be run by itself"""
 
+    _callstack_regex = re.compile(r'  File "(\/.[^"]*)", line (\d*),(.*)')
+
     def __init__(self, environment):
         super().__init__(environment)
         if self.host is None:
@@ -351,6 +356,56 @@ def __init__(self, environment):
         The client support cookies, and therefore keeps the session between HTTP requests.
         """
 
+    @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
+
 
 class FastRequest(CompatRequest):
     payload: Optional[str] = None
@@ -572,3 +627,8 @@ def failure(self, exc):
         if not isinstance(exc, Exception):
             exc = CatchResponseError(exc)
         self._manual_result = exc
+
+
+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

tox.ini

@@ -28,7 +28,7 @@ allowlist_externals =
 commands =
     python3 -m pip install .
     python3 -m coverage run -m unittest discover []
-    bash -ec "PYTHONUNBUFFERED=1 timeout 10s python3 examples/rest_ex.py >out.txt 2>err.txt || true"
+    bash -ec "PYTHONUNBUFFERED=1 timeout 10s python3 examples/rest.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