locust 2.29.1.dev34__py3-none-any.whl → 2.29.1.dev39__py3-none-any.whl

locust/_version.py

@@ -3,6 +3,7 @@
 TYPE_CHECKING = False
 if TYPE_CHECKING:
     from typing import Tuple, Union
+
     VERSION_TUPLE = Tuple[Union[int, str], ...]
 else:
     VERSION_TUPLE = object
@@ -12,5 +13,8 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '2.29.1.dev34'
-__version_tuple__ = version_tuple = (2, 29, 1, 'dev34')
+
+__version__ = "2.29.1.dev39"
+version = __version__
+__version_tuple__ = (2, 29, 1, "dev39")
+version_tuple = __version_tuple__

locust/contrib/fasthttp.py

@@ -12,12 +12,11 @@ import socket
 import time
 import traceback
 from base64 import b64encode
-from collections.abc import Generator
 from contextlib import contextmanager
 from http.cookiejar import CookieJar
 from json.decoder import JSONDecodeError
 from ssl import SSLError
-from typing import Any, Callable, cast
+from typing import TYPE_CHECKING, cast
 from urllib.parse import urlparse, urlunparse
 
 import gevent
@@ -32,6 +31,36 @@ from geventhttpclient.useragent import CompatRequest, CompatResponse, Connection
 # borrow requests's content-type header parsing
 from requests.utils import get_encoding_from_headers
 
+if TYPE_CHECKING:
+    import sys
+    from collections.abc import Callable, Generator
+    from typing import TypedDict
+
+    if sys.version_info >= (3, 11):
+        from typing import Unpack
+    else:
+        from typing_extensions import Unpack
+
+    class PostKwargs(TypedDict, total=False):
+        name: str | None
+        catch_response: bool
+        stream: bool
+        headers: dict | None
+        auth: tuple[str | bytes, str | bytes] | None
+        allow_redirects: bool
+        context: dict
+
+    class PutKwargs(PostKwargs, total=False):
+        json: dict | None
+
+    class PatchKwargs(PostKwargs, total=False):
+        json: dict | None
+
+    class RESTKwargs(PostKwargs, total=False):
+        data: str | dict | None
+        json: dict | None
+
+
 # Monkey patch geventhttpclient.useragent.CompatRequest so that Cookiejar works with Python >= 3.3
 # More info: https://github.com/requests/requests/pull/871
 CompatRequest.unverifiable = False
@@ -85,7 +114,7 @@ class FastHttpSession:
         client_pool: HTTPClientPool | None = None,
         ssl_context_factory: Callable | None = None,
         **kwargs,
-    ):
+    ) -> None:
         self.environment = environment
         self.base_url = base_url
         self.cookiejar = CookieJar()
@@ -117,14 +146,14 @@ class FastHttpSession:
             # store authentication header (we construct this by using _basic_auth_str() function from requests.auth)
             self.auth_header = _construct_basic_auth_str(parsed_url.username, parsed_url.password)
 
-    def _build_url(self, path):
+    def _build_url(self, path: str) -> str:
         """prepend url with hostname unless it's already an absolute URL"""
         if absolute_http_url_regexp.match(path):
             return path
         else:
             return f"{self.base_url}{path}"
 
-    def _send_request_safe_mode(self, method, url, **kwargs):
+    def _send_request_safe_mode(self, method: str, url: str, **kwargs):
         """
         Send an HTTP request, and catch any exception that might occur due to either
         connection problems, or invalid HTTP status codes
@@ -155,9 +184,9 @@ class FastHttpSession:
         catch_response: bool = False,
         stream: bool = False,
         headers: dict | None = None,
-        auth=None,
+        auth: tuple[str | bytes, str | bytes] | None = None,
         json: dict | None = None,
-        allow_redirects=True,
+        allow_redirects: bool = True,
         context: dict = {},
         **kwargs,
     ) -> ResponseContextManager | FastResponse:
@@ -187,6 +216,7 @@ class FastHttpSession:
             and can instead be consumed by accessing the stream attribute on the Response object.
             Another side effect of setting stream to True is that the time for downloading the response
             content will not be accounted for in the request time that is reported by Locust.
+        :param allow_redirects: (optional) Set to True by default.
         """
         # prepend url with hostname unless it's already an absolute URL
         built_url = self._build_url(url)
@@ -250,7 +280,7 @@ class FastHttpSession:
         # Record the consumed time
         # Note: This is intentionally placed after we record the content_size above, since
         # we'll then trigger fetching of the body (unless stream=True)
-        request_meta["response_time"] = int((time.perf_counter() - start_perf_counter) * 1000)
+        request_meta["response_time"] = (time.perf_counter() - start_perf_counter) * 1000
 
         if catch_response:
             return ResponseContextManager(response, environment=self.environment, request_meta=request_meta)
@@ -263,30 +293,37 @@ class FastHttpSession:
             self.environment.events.request.fire(**request_meta)
             return response
 
-    def delete(self, url, **kwargs):
+    def delete(self, url: str, **kwargs: Unpack[RESTKwargs]) -> ResponseContextManager | FastResponse:
+        """Sends a DELETE request"""
         return self.request("DELETE", url, **kwargs)
 
-    def get(self, url, **kwargs):
+    def get(self, url: str, **kwargs: Unpack[RESTKwargs]) -> ResponseContextManager | FastResponse:
         """Sends a GET request"""
         return self.request("GET", url, **kwargs)
 
-    def head(self, url, **kwargs):
+    def head(self, url: str, **kwargs: Unpack[RESTKwargs]) -> ResponseContextManager | FastResponse:
         """Sends a HEAD request"""
         return self.request("HEAD", url, **kwargs)
 
-    def options(self, url, **kwargs):
+    def options(self, url: str, **kwargs: Unpack[RESTKwargs]) -> ResponseContextManager | FastResponse:
         """Sends a OPTIONS request"""
         return self.request("OPTIONS", url, **kwargs)
 
-    def patch(self, url, data=None, **kwargs):
-        """Sends a POST request"""
+    def patch(
+        self, url: str, data: str | dict | None = None, **kwargs: Unpack[PatchKwargs]
+    ) -> ResponseContextManager | FastResponse:
+        """Sends a PATCH request"""
         return self.request("PATCH", url, data=data, **kwargs)
 
-    def post(self, url, data=None, **kwargs):
+    def post(
+        self, url: str, data: str | dict | None = None, json: dict | None = None, **kwargs: Unpack[PostKwargs]
+    ) -> ResponseContextManager | FastResponse:
         """Sends a POST request"""
-        return self.request("POST", url, data=data, **kwargs)
+        return self.request("POST", url, data=data, json=json, **kwargs)
 
-    def put(self, url, data=None, **kwargs):
+    def put(
+        self, url: str, data: str | dict | None = None, **kwargs: Unpack[PutKwargs]
+    ) -> ResponseContextManager | FastResponse:
         """Sends a PUT request"""
         return self.request("PUT", url, data=data, **kwargs)
 
@@ -446,7 +483,7 @@ class FastResponse(CompatResponse):
 
     _response: HTTPSocketPoolResponse | None = None
 
-    encoding: str | None = None
+    encoding: str | float | None = None
     """In some cases setting the encoding explicitly is needed. If so, do it before calling .text"""
 
     request: FastRequest | None = None

locust/dispatch.py

@@ -90,13 +90,13 @@ class UsersDispatcher(Iterator):
         assert len(user_classes) > 0
         assert len(set(self._user_classes)) == len(self._user_classes)
 
-        self._target_user_count: int = None
+        self._target_user_count: int = 0
 
-        self._spawn_rate: float = None
+        self._spawn_rate: float = 0.0
 
-        self._user_count_per_dispatch_iteration: int = None
+        self._user_count_per_dispatch_iteration: int = 0
 
-        self._wait_between_dispatch: float = None
+        self._wait_between_dispatch: float = 0.0
 
         self._initial_users_on_workers = {
             worker_node.id: {user_class.__name__: 0 for user_class in self._user_classes}
@@ -107,7 +107,8 @@ class UsersDispatcher(Iterator):
 
         self._current_user_count = self.get_current_user_count()
 
-        self._dispatcher_generator: Generator[dict[str, dict[str, int]], None, None] = None
+        self._dispatcher_generator: Generator[dict[str, dict[str, int]], None, None] = None  # type: ignore
+        # a generator is assigned (in new_dispatch()) to _dispatcher_generator before it's used
 
         self._user_generator = self._user_gen()
 
@@ -131,7 +132,7 @@ class UsersDispatcher(Iterator):
         return sum(map(sum, map(dict.values, self._users_on_workers.values())))
 
     @property
-    def dispatch_in_progress(self):
+    def dispatch_in_progress(self) -> bool:
         return self._dispatch_in_progress
 
     @property

locust/event.py

@@ -235,6 +235,38 @@ class Events:
     Fired when the CPU usage exceeds runners.CPU_WARNING_THRESHOLD (90% by default)
     """
 
+    heartbeat_sent: EventHook
+    """
+    Fired when a heartbeat is sent by master to a worker.
+
+    Event arguments:
+
+    :param client_id: worker client id
+    :param timestamp: time in seconds since the epoch (float) when the event occured
+    """
+
+    heartbeat_received: EventHook
+    """
+    Fired when a heartbeat is received by a worker from master.
+
+    Event arguments:
+
+    :param client_id: worker client id
+    :param timestamp: time in seconds since the epoch (float) when the event occured
+    """
+
+    usage_monitor: EventHook
+    """
+    Fired every runners.CPU_MONITOR_INTERVAL (5.0 seconds by default) with information about
+    current CPU and memory usage.
+
+    Event arguments:
+
+    :param environment: locust environment
+    :param cpu_usage: current CPU usage in percent
+    :param memory_usage: current memory usage (RSS) in bytes
+    """
+
     def __init__(self):
         # For backward compatibility use also values of class attributes
         for name, value in vars(type(self)).items():

locust/runners.py

@@ -15,19 +15,9 @@ import traceback
 from abc import abstractmethod
 from collections import defaultdict
 from collections.abc import Iterator, MutableMapping, ValuesView
-from operator import (
-    itemgetter,
-    methodcaller,
-)
+from operator import itemgetter, methodcaller
 from types import TracebackType
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    NoReturn,
-    TypedDict,
-    cast,
-)
+from typing import TYPE_CHECKING, Any, Callable, NoReturn, TypedDict, cast
 from uuid import uuid4
 
 import gevent
@@ -40,15 +30,8 @@ from . import argument_parser
 from .dispatch import UsersDispatcher
 from .exception import RPCError, RPCReceiveError, RPCSendError
 from .log import get_logs, greenlet_exception_logger
-from .rpc import (
-    Message,
-    rpc,
-)
-from .stats import (
-    RequestStats,
-    StatsError,
-    setup_distributed_stats_event_listeners,
-)
+from .rpc import Message, rpc
+from .stats import RequestStats, StatsError, setup_distributed_stats_event_listeners
 
 if TYPE_CHECKING:
     from . import User
@@ -106,7 +89,7 @@ class Runner:
         self.spawning_greenlet: gevent.Greenlet | None = None
         self.shape_greenlet: gevent.Greenlet | None = None
         self.shape_last_tick: tuple[int, float] | tuple[int, float, list[type[User]] | None] | None = None
-        self.current_cpu_usage: int = 0
+        self.current_cpu_usage: float = 0.0
         self.cpu_warning_emitted: bool = False
         self.worker_cpu_warning_emitted: bool = False
         self.current_memory_usage: int = 0
@@ -308,6 +291,10 @@ class Runner:
                         f"CPU usage above {CPU_WARNING_THRESHOLD}%! This may constrain your throughput and may even give inconsistent response time measurements! See https://docs.locust.io/en/stable/running-distributed.html for how to distribute the load over multiple CPU cores or machines"
                     )
                     self.cpu_warning_emitted = True
+
+            self.environment.events.usage_monitor.fire(
+                environment=self.environment, cpu_usage=self.current_cpu_usage, memory_usage=self.current_memory_usage
+            )
             gevent.sleep(CPU_MONITOR_INTERVAL)
 
     @abstractmethod
@@ -1102,6 +1089,7 @@ class MasterRunner(DistributedRunner):
                         )
                     if "current_memory_usage" in msg.data:
                         c.memory_usage = msg.data["current_memory_usage"]
+                    self.environment.events.heartbeat_sent.fire(client_id=msg.node_id, timestamp=time.time())
                     self.server.send_to_client(Message("heartbeat", None, msg.node_id))
                 else:
                     logging.debug(f"Got heartbeat message from unknown worker {msg.node_id}")
@@ -1399,6 +1387,9 @@ class WorkerRunner(DistributedRunner):
                 self.reset_connection()
             elif msg.type == "heartbeat":
                 self.last_heartbeat_timestamp = time.time()
+                self.environment.events.heartbeat_received.fire(
+                    client_id=msg.node_id, timestamp=self.last_heartbeat_timestamp
+                )
             elif msg.type == "update_user_class":
                 self.environment.update_user_class(msg.data)
             elif msg.type == "spawning_complete":

locust/stats.py

@@ -9,25 +9,12 @@ import signal
 import time
 from abc import abstractmethod
 from collections import OrderedDict, defaultdict, namedtuple
-from collections import (
-    OrderedDict as OrderedDictType,
-)
 from collections.abc import Iterable
 from copy import copy
 from html import escape
 from itertools import chain
-from tempfile import NamedTemporaryFile
 from types import FrameType
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    NoReturn,
-    Protocol,
-    TypedDict,
-    TypeVar,
-    cast,
-)
+from typing import TYPE_CHECKING, Any, Callable, NoReturn, Protocol, TypedDict, TypeVar, cast
 
 import gevent
 
@@ -319,12 +306,12 @@ class StatsEntry:
         A {response_time => count} dict that holds the response time distribution of all
         the requests.
 
-        The keys (the response time in ms) are rounded to store 1, 2, ... 9, 10, 20. .. 90,
-        100, 200 .. 900, 1000, 2000 ... 9000, in order to save memory.
+        The keys (the response time in ms) are rounded to store 1, 2, ... 98, 99, 100, 110, 120, ... 980, 990, 1000,
+        1100, 1200, ... 9800, 9900, 10_000, 11_000, 12_000 ... in order to save memory.
 
         This dict is used to calculate the median and percentile response times.
         """
-        self.response_times_cache: OrderedDictType[int, CachedResponseTimes] | None = None
+        self.response_times_cache: OrderedDict[int, CachedResponseTimes] | None = None
         """
         If use_response_times_cache is set to True, this will be a {timestamp => CachedResponseTimes()}
         OrderedDict that holds a copy of the response_times dict for each of the last 20 seconds.

locust/user/sequential_taskset.py

@@ -1,6 +1,6 @@
 from locust.exception import LocustError
 
-import logging
+from itertools import cycle
 
 from .task import TaskSet, TaskSetMeta
 
@@ -26,8 +26,12 @@ class SequentialTaskSetMeta(TaskSetMeta):
                 # compared to methods declared with @task
                 if isinstance(value, list):
                     new_tasks.extend(value)
+                elif isinstance(value, dict):
+                    for task, weight in value.items():
+                        for _ in range(weight):
+                            new_tasks.append(task)
                 else:
-                    raise ValueError("On SequentialTaskSet the task attribute can only be set to a list")
+                    raise ValueError("The 'tasks' attribute can only be set to list or dict")
 
             if "locust_task_weight" in dir(value):
                 # method decorated with @task
@@ -52,13 +56,11 @@ class SequentialTaskSet(TaskSet, metaclass=SequentialTaskSetMeta):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self._task_index = 0
+        self._task_cycle = cycle(self.tasks)
 
     def get_next_task(self):
         if not self.tasks:
             raise LocustError(
                 "No tasks defined. Use the @task decorator or set the 'tasks' attribute of the SequentialTaskSet"
             )
-        task = self.tasks[self._task_index % len(self.tasks)]
-        self._task_index += 1
-        return task
+        return next(self._task_cycle)