locust 2.29.1.dev31__tar.gz → 2.29.1.dev39__tar.gz

{locust-2.29.1.dev31 → locust-2.29.1.dev39}/PKG-INFO

@@ -1,15 +1,15 @@
 Metadata-Version: 2.1
 Name: locust
-Version: 2.29.1.dev31
+Version: 2.29.1.dev39
 Summary: Developer-friendly load testing framework
+Home-page: https://locust.io/
 License: MIT
-Project-URL: Homepage, https://github.com/locustio/locust
-Project-URL: Documentation, https://docs.locust.io/
-Project-URL: Code, https://github.com/locustio/locust
-Project-URL: Help/Questions, https://stackoverflow.com/questions/tagged/locust
-Project-URL: Issue tracker, https://github.com/locustio/locust/issues
-Classifier: Topic :: Software Development :: Testing :: Traffic Generation
+Author: Jonatan Heyman
+Maintainer: Lars Holmberg
+Requires-Python: >=3.9
 Classifier: Development Status :: 5 - Production/Stable
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: System Administrators
 Classifier: License :: OSI Approved :: MIT License
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python
@@ -18,29 +18,32 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
-Classifier: Intended Audience :: Developers
-Classifier: Intended Audience :: System Administrators
 Classifier: Topic :: Software Development :: Testing
 Classifier: Topic :: Software Development :: Testing :: Traffic Generation
 Classifier: Topic :: System :: Distributed Computing
-Requires-Python: >=3.9
+Requires-Dist: ConfigArgParse (>=1.5.5)
+Requires-Dist: Flask-Cors (>=3.0.10)
+Requires-Dist: Flask-Login (>=0.6.3)
+Requires-Dist: Werkzeug (>=2.0.0)
+Requires-Dist: flask (>=2.0.0)
+Requires-Dist: gevent (>=22.10.2)
+Requires-Dist: geventhttpclient (>=2.3.1)
+Requires-Dist: msgpack (>=1.0.0)
+Requires-Dist: psutil (>=5.9.1)
+Requires-Dist: pywin32 ; sys_platform == "win32"
+Requires-Dist: pyzmq (>=25.0.0)
+Requires-Dist: requests (>=2.26.0) ; python_full_version <= "3.11.0"
+Requires-Dist: requests (>=2.32.2) ; python_full_version > "3.11.0"
+Requires-Dist: tomli (>=1.1.0) ; python_version < "3.11"
+Requires-Dist: typing_extensions (>=4.6.0) ; python_version < "3.11"
+Project-URL: Documentation, https://docs.locust.io/
+Project-URL: Help/Questions, https://github.com/orgs/locustio/discussions/
+Project-URL: Issue Tracker, https://github.com/locustio/locust/issues
+Project-URL: Repository, https://github.com/locustio/locust
+Project-URL: Stack Overflow, https://stackoverflow.com/questions/tagged/locust
+Project-URL: Slack, https://locustio.slack.com/
+Project-URL: Slack/Signup, https://communityinviter.com/apps/locustio/locust
 Description-Content-Type: text/markdown
-License-File: LICENSE
-Requires-Dist: gevent>=22.10.2
-Requires-Dist: flask>=2.0.0
-Requires-Dist: Werkzeug>=2.0.0
-Requires-Dist: requests>=2.26.0
-Requires-Dist: requests>=2.32.2; python_version > "3.11"
-Requires-Dist: msgpack>=1.0.0
-Requires-Dist: pyzmq>=25.0.0
-Requires-Dist: geventhttpclient>=2.3.1
-Requires-Dist: ConfigArgParse>=1.5.5
-Requires-Dist: tomli>=1.1.0; python_version < "3.11"
-Requires-Dist: typing_extensions>=4.6.0; python_version < "3.11"
-Requires-Dist: psutil>=5.9.1
-Requires-Dist: Flask-Login>=0.6.3
-Requires-Dist: Flask-Cors>=3.0.10
-Requires-Dist: pywin32; platform_system == "Windows"
 
 # Locust
 
@@ -112,6 +115,7 @@ Locust's code base is intentionally kept small and doesn't solve everything out
 
 * Documentation: [docs.locust.io](https://docs.locust.io)
 * Support/Questions: [StackOverflow](https://stackoverflow.com/questions/tagged/locust)
+* Github Discussions: [Github Discussions](https://github.com/orgs/locustio/discussions)
 * Chat/discussion: [Slack](https://locustio.slack.com) [(signup)](https://communityinviter.com/apps/locustio/locust)
 
 ## Authors
@@ -124,3 +128,4 @@ Locust's code base is intentionally kept small and doesn't solve everything out
 ## License
 
 Open source licensed under the MIT license (see _LICENSE_ file for details).
+

{locust-2.29.1.dev31 → locust-2.29.1.dev39}/README.md

@@ -68,6 +68,7 @@ Locust's code base is intentionally kept small and doesn't solve everything out
 
 * Documentation: [docs.locust.io](https://docs.locust.io)
 * Support/Questions: [StackOverflow](https://stackoverflow.com/questions/tagged/locust)
+* Github Discussions: [Github Discussions](https://github.com/orgs/locustio/discussions)
 * Chat/discussion: [Slack](https://locustio.slack.com) [(signup)](https://communityinviter.com/apps/locustio/locust)
 
 ## Authors

{locust-2.29.1.dev31 → locust-2.29.1.dev39}/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.dev31'
-__version_tuple__ = version_tuple = (2, 29, 1, 'dev31')
+
+__version__ = "2.29.1.dev39"
+version = __version__
+__version_tuple__ = (2, 29, 1, "dev39")
+version_tuple = __version_tuple__

{locust-2.29.1.dev31 → locust-2.29.1.dev39}/locust/clients.py

@@ -7,7 +7,7 @@ from typing import TYPE_CHECKING
 from urllib.parse import urlparse, urlunparse
 
 import requests
-from requests import Request, Response
+from requests import Response
 from requests.adapters import HTTPAdapter
 from requests.auth import HTTPBasicAuth
 from requests.exceptions import InvalidSchema, InvalidURL, MissingSchema, RequestException
@@ -54,8 +54,10 @@ absolute_http_url_regexp = re.compile(r"^https?://", re.I)
 
 
 class LocustResponse(Response):
-    def raise_for_status(self):
-        if hasattr(self, "error") and self.error:
+    error: Exception | None = None
+
+    def raise_for_status(self) -> None:
+        if self.error:
             raise self.error
         Response.raise_for_status(self)
 
@@ -112,7 +114,7 @@ class HttpSession(requests.Session):
         self.mount("https://", LocustHttpAdapter(pool_manager=pool_manager))
         self.mount("http://", LocustHttpAdapter(pool_manager=pool_manager))
 
-    def _build_url(self, path):
+    def _build_url(self, path) -> str:
         """prepend url with hostname unless it's already an absolute URL"""
         if absolute_http_url_regexp.match(path):
             return path
@@ -191,7 +193,7 @@ class HttpSession(requests.Session):
         response_time = (time.perf_counter() - start_perf_counter) * 1000
 
         request_before_redirect = (response.history and response.history[0] or response).request
-        url = request_before_redirect.url
+        url = request_before_redirect.url  # type: ignore
 
         if not name:
             name = request_before_redirect.path_url
@@ -225,7 +227,7 @@ class HttpSession(requests.Session):
                 pass
             return response
 
-    def _send_request_safe_mode(self, method, url, **kwargs):
+    def _send_request_safe_mode(self, method, url, **kwargs) -> Response | LocustResponse:
         """
         Send an HTTP request, and catch any exception that might occur due to connection problems.
 
@@ -238,8 +240,8 @@ class HttpSession(requests.Session):
         except RequestException as e:
             r = LocustResponse()
             r.error = e
-            r.status_code = 0  # with this status_code, content returns None
-            r.request = Request(method, url).prepare()
+            r.status_code = 0
+            r.request = e.request  # type: ignore
             return r
 
     def get(
@@ -440,17 +442,11 @@ class LocustHttpAdapter(HTTPAdapter):
 
 
 # Monkey patch Response class to give some guidance
-def _success(self):
-    raise LocustError(
-        "If you want to change the state of the request, you must pass catch_response=True. See http://docs.locust.io/en/stable/writing-a-locustfile.html#validating-responses"
-    )
-
-
-def _failure(self):
+def _missing_catch_response_True(self, *_args, **_kwargs):
     raise LocustError(
-        "If you want to change the state of the request, you must pass catch_response=True. See http://docs.locust.io/en/stable/writing-a-locustfile.html#validating-responses"
+        "If you want to change the state of the request using .success() or .failure(), you must pass catch_response=True. See http://docs.locust.io/en/stable/writing-a-locustfile.html#validating-responses"
     )
 
 
-Response.success = _success  # type: ignore[attr-defined]
-Response.failure = _failure  # type: ignore[attr-defined]
+Response.success = _missing_catch_response_True  # type: ignore[attr-defined]
+Response.failure = _missing_catch_response_True  # type: ignore[attr-defined]

{locust-2.29.1.dev31 → locust-2.29.1.dev39}/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)
 
@@ -324,6 +361,12 @@ class FastHttpUser(User):
     Note that setting this value has no effect when custom client_pool was given, and you need to spawn a your own gevent pool
     to use it (as Users only have one greenlet). See test_fasthttp.py / test_client_pool_concurrency for an example."""
 
+    proxy_host: str | None = None
+    """Parameter passed to FastHttpSession"""
+
+    proxy_port: int | None = None
+    """Parameter passed to FastHttpSession"""
+
     client_pool: HTTPClientPool | None = None
     """HTTP client pool to use. If not given, a new pool is created per single user."""
 
@@ -355,6 +398,8 @@ class FastHttpUser(User):
             client_pool=self.client_pool,
             ssl_context_factory=self.ssl_context_factory,
             headers=self.default_headers,
+            proxy_host=self.proxy_host,
+            proxy_port=self.proxy_port,
         )
         """
         Instance of HttpSession that is created upon instantiation of User.
@@ -438,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-2.29.1.dev31 → locust-2.29.1.dev39}/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-2.29.1.dev31 → locust-2.29.1.dev39}/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-2.29.1.dev31 → locust-2.29.1.dev39}/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-2.29.1.dev31 → locust-2.29.1.dev39}/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-2.29.1.dev31 → locust-2.29.1.dev39}/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)