locust 2.46.5.dev8__tar.gz → 2.46.6.dev4__tar.gz

{locust-2.46.5.dev8 → locust-2.46.6.dev4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: locust
-Version: 2.46.5.dev8
+Version: 2.46.6.dev4
 Summary: Developer-friendly load testing framework
 Project-URL: Homepage, https://locust.cloud/
 Project-URL: Repository, https://github.com/locustio/locust

{locust-2.46.5.dev8 → locust-2.46.6.dev4}/locust/__init__.py

@@ -31,7 +31,7 @@ from .user.markov_taskset import MarkovTaskSet, transition, transitions
 from .user.sequential_taskset import SequentialTaskSet
 from .user.task import TaskSet, tag, task
 from .user.users import HttpUser, User
-from .user.wait_time import between, constant, constant_pacing, constant_throughput
+from .user.wait_time import between, constant, constant_pacing, constant_throughput, poisson
 
 events = Events()
 
@@ -51,6 +51,7 @@ __all__ = (
     "constant",
     "constant_pacing",
     "constant_throughput",
+    "poisson",
     "events",
     "LoadTestShape",
     "run_single_user",

{locust-2.46.5.dev8 → locust-2.46.6.dev4}/locust/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '2.46.5.dev8'
-__version_tuple__ = version_tuple = (2, 46, 5, 'dev8')
+__version__ = version = '2.46.6.dev4'
+__version_tuple__ = version_tuple = (2, 46, 6, 'dev4')
 
 __commit_id__ = commit_id = None

{locust-2.46.5.dev8 → locust-2.46.6.dev4}/locust/clients.py

@@ -478,6 +478,10 @@ class ResponseContextManager(Response):
         self._manual_result = exc
 
 
+class RestResponseContextManager(ResponseContextManager):
+    js: dict  # This is technically an Optional, but I dont want to force everyone to check it
+
+
 class LocustHttpAdapter(HTTPAdapter):
     def __init__(self, pool_manager: PoolManager | None, *args, **kwargs):
         self.poolmanager = pool_manager

{locust-2.46.5.dev8 → locust-2.46.6.dev4}/locust/contrib/fasthttp.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 from locust.exception import CatchResponseError, LocustError, ResponseError, StopTest
 from locust.user import User
+from locust.user.users import RestMixin
 from locust.util.deprecation import DeprecatedFastHttpLocustClass as FastHttpLocust  # noqa: F401
 
 import json
@@ -9,14 +10,11 @@ import json as unshadowed_json  # some methods take a named parameter called jso
 import re
 import socket
 import time
-import traceback
 import zlib
 from base64 import b64encode
-from contextlib import contextmanager
 from http.cookiejar import CookieJar
-from json.decoder import JSONDecodeError
 from ssl import SSLError
-from typing import TYPE_CHECKING, cast
+from typing import TYPE_CHECKING
 from urllib.parse import urlparse, urlunparse
 
 import gevent
@@ -355,7 +353,7 @@ class FastHttpSession:
         return self.request("PUT", url, data=data, **kwargs)
 
 
-class FastHttpUser(User):
+class FastHttpUser(User, RestMixin["RestResponseContextManager"]):
     """
     FastHttpUser provides the same API as HttpUser, but uses geventhttpclient instead of python-requests
     as its underlying client. It uses considerably less CPU on the load generator, and should work
@@ -413,8 +411,6 @@ 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) -> None:
         super().__init__(environment)
         if self.host is None:
@@ -443,66 +439,6 @@ class FastHttpUser(User):
         The client support cookies, and therefore keeps the session between HTTP requests.
         """
 
-    @contextmanager
-    def rest(self, method, url, headers: dict | None = None, **kwargs) -> Generator[RestResponseContextManager]:
-        """
-        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)
-        """
-        headers = headers or {}
-        if not ("Content-Type" in headers or "content-type" in headers):
-            headers["Content-Type"] = "application/json"
-        if not ("Accept" in headers or "accept" in headers):
-            headers["Accept"] = "application/json"
-        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.content is None:
-                resp.failure(str(resp.error))
-            elif 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"):
-                    if m := self._callstack_regex.match(l):
-                        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}")
-
-    @contextmanager
-    def rest_(self, method, url, name=None, **kwargs) -> Generator[RestResponseContextManager]:
-        """
-        Some REST api:s use a timestamp as part of their query string (mainly to break through caches).
-        This is a convenience method for that, appending a _=<timestamp> parameter automatically
-        """
-        separator = "&" if "?" in url else "?"
-        if name is None:
-            name = url + separator + "_=..."
-        with self.rest(method, f"{url}{separator}_={int(time.time() * 1000)}", name=name, **kwargs) as resp:
-            yield resp
-
 
 class FastRequest(CompatRequest):
     payload: str | None = None

{locust-2.46.5.dev8 → locust-2.46.6.dev4}/locust/user/users.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from locust.clients import HttpSession
+from locust.clients import HttpSession, RestResponseContextManager
 from locust.exception import CatchResponseError, StopTest, StopUser
 from locust.user.task import (
     LOCUST_STATE_RUNNING,
@@ -14,11 +14,14 @@ from locust.user.wait_time import constant
 from locust.util import deprecation
 
 import logging
+import re
 import sys
 import time
 import traceback
 from collections.abc import Callable
-from typing import TYPE_CHECKING, final
+from contextlib import contextmanager
+from json.decoder import JSONDecodeError
+from typing import TYPE_CHECKING, Any, Generic, TypeVar, final
 
 from gevent import GreenletExit, greenlet
 from gevent.pool import Group
@@ -32,6 +35,8 @@ else:
     from typing_extensions import override
 
 if TYPE_CHECKING:
+    from collections.abc import Generator
+
     import pytest
 
 logger = logging.getLogger(__name__)
@@ -250,7 +255,78 @@ class User(metaclass=UserMeta):
         return ".".join(filter(lambda x: x != "<locals>", (cls.__module__ + "." + cls.__qualname__).split(".")))
 
 
-class HttpUser(User):
+RestResponseT = TypeVar("RestResponseT")
+
+
+class RestMixin(Generic[RestResponseT]):
+    """
+    Provides the ``rest`` and ``rest_`` methods for User classes that have an HTTP ``client`` (HttpUser and FastHttpUser)
+    """
+
+    _callstack_regex = re.compile(r'  File "(\/.[^"]*)", line (\d*),(.*)')
+
+    @contextmanager
+    def rest(self, method, url, headers: dict | None = None, **kwargs) -> Generator[RestResponseT]:
+        """
+        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)
+        """
+        headers = headers or {}
+        if not ("Content-Type" in headers or "content-type" in headers):
+            headers["Content-Type"] = "application/json"
+        if not ("Accept" in headers or "accept" in headers):
+            headers["Accept"] = "application/json"
+        with self.client.request(method, url, catch_response=True, headers=headers, **kwargs) as r:  # type: ignore[attr-defined]
+            resp: Any = r
+            resp.js = None
+            if resp.content is None:
+                resp.failure(str(resp.error))
+            elif 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"):
+                    if m := self._callstack_regex.match(l):
+                        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}")
+
+    @contextmanager
+    def rest_(self, method, url, name=None, **kwargs) -> Generator[RestResponseT]:
+        """
+        Some REST api:s use a timestamp as part of their query string (mainly to break through caches).
+        This is a convenience method for that, appending a _=<timestamp> parameter automatically
+        """
+        separator = "&" if "?" in url else "?"
+        if name is None:
+            name = url + separator + "_=..."
+        with self.rest(method, f"{url}{separator}_={int(time.time() * 1000)}", name=name, **kwargs) as resp:
+            yield resp
+
+
+class HttpUser(User, RestMixin[RestResponseContextManager]):
     """
     Represents an HTTP "user" which is to be spawned and attack the system that is to be load tested.
 

{locust-2.46.5.dev8 → locust-2.46.6.dev4}/locust/user/wait_time.py

@@ -84,3 +84,30 @@ def constant_throughput(task_runs_per_second: float) -> Callable[["User"], float
     the next task.
     """
     return constant_pacing(1 / task_runs_per_second)
+
+
+def poisson(rate: float) -> Callable[["User"], float]:
+    """
+    Returns a function that will return a random wait time sampled from an exponential
+    distribution with mean ``1/rate`` seconds. Tasks will therefore arrive according to a
+    Poisson process with mean rate ``rate`` tasks per second, which is a good approximation
+    of how real users arrive.
+
+    The average number of task runs per second will be ``rate``, but unlike
+    :py:func:`constant_throughput <locust.wait_time.constant_throughput>` the individual
+    intervals vary randomly, so bursts and quiet periods occur naturally.
+
+    Example::
+
+        class MyUser(User):
+            wait_time = poisson(2)  # on average, each user runs 2 tasks per second
+            @task
+            def my_task(self):
+                ...
+
+    If you have multiple requests in a task your RPS will of course be higher than the
+    specified throughput.
+    """
+    if rate <= 0:
+        raise ValueError(f"poisson() requires a positive rate, got: {rate}")
+    return lambda instance: random.expovariate(rate)