locust 2.23.2.dev7__py3-none-any.whl → 2.23.2.dev18__py3-none-any.whl

locust/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '2.23.2.dev7'
-__version_tuple__ = version_tuple = (2, 23, 2, 'dev7')
+__version__ = version = '2.23.2.dev18'
+__version_tuple__ = version_tuple = (2, 23, 2, 'dev18')

locust/dispatch.py

@@ -5,9 +5,8 @@ import itertools
 import math
 import time
 from collections import defaultdict
-from collections.abc import Iterator
 from operator import attrgetter
-from typing import TYPE_CHECKING, Generator
+from typing import TYPE_CHECKING, Generator, Iterator
 
 import gevent
 from roundrobin import smooth
@@ -164,7 +163,9 @@ class UsersDispatcher(Iterator):
 
         self._dispatch_in_progress = False
 
-    def new_dispatch(self, target_user_count: int, spawn_rate: float, user_classes: list | None = None) -> None:
+    def new_dispatch(
+        self, target_user_count: int, spawn_rate: float, user_classes: list[type[User]] | None = None
+    ) -> None:
         """
         Initialize a new dispatch cycle.
 

locust/env.py

@@ -5,6 +5,7 @@ from typing import Callable, TypeVar
 
 from configargparse import Namespace
 
+from .dispatch import UsersDispatcher
 from .event import Events
 from .exception import RunnerAlreadyExistsError
 from .runners import LocalRunner, MasterRunner, Runner, WorkerRunner
@@ -35,6 +36,7 @@ class Environment:
         available_user_classes: dict[str, User] | None = None,
         available_shape_classes: dict[str, LoadTestShape] | None = None,
         available_user_tasks: dict[str, list[TaskSet | Callable]] | None = None,
+        dispatcher_class: type[UsersDispatcher] = UsersDispatcher,
     ):
         self.runner: Runner | None = None
         """Reference to the :class:`Runner <locust.runners.Runner>` instance"""
@@ -95,6 +97,8 @@ class Environment:
         """List of the available Shape Classes to pick from in the ShapeClass Picker"""
         self.available_user_tasks = available_user_tasks
         """List of the available Tasks per User Classes to pick from in the Task Picker"""
+        self.dispatcher_class = dispatcher_class
+        """A user dispatcher class that decides how users are spawned, default :class:`UsersDispatcher <locust.dispatch.UsersDispatcher>`"""
 
         self._remove_user_classes_with_weight_zero()
         self._validate_user_class_name_uniqueness()

locust/runners.py

@@ -492,13 +492,13 @@ class LocalRunner(Runner):
             self.update_state(STATE_SPAWNING)
 
         if self._users_dispatcher is None:
-            self._users_dispatcher = UsersDispatcher(
+            self._users_dispatcher = self.environment.dispatcher_class(
                 worker_nodes=[self._local_worker_node], user_classes=self.user_classes
             )
 
         logger.info("Ramping to %d users at a rate of %.2f per second" % (user_count, spawn_rate))
 
-        cast(UsersDispatcher, self._users_dispatcher).new_dispatch(user_count, spawn_rate, user_classes)
+        self._users_dispatcher.new_dispatch(user_count, spawn_rate, user_classes)
 
         try:
             for dispatched_users in self._users_dispatcher:
@@ -750,7 +750,7 @@ class MasterRunner(DistributedRunner):
         self.spawn_rate = spawn_rate
 
         if self._users_dispatcher is None:
-            self._users_dispatcher = UsersDispatcher(
+            self._users_dispatcher = self.environment.dispatcher_class(
                 worker_nodes=list(self.clients.values()), user_classes=self.user_classes
             )
 

locust/test/test_env.py

@@ -1,6 +1,7 @@
 from locust import (
     constant,
 )
+from locust.dispatch import UsersDispatcher
 from locust.env import Environment, LoadTestShape
 from locust.user import (
     User,
@@ -201,6 +202,18 @@ class TestEnvironment(LocustTestCase):
         ):
             Environment(user_classes=[MyUserWithSameName1], shape_class=SubLoadTestShape)
 
+    def test_dispatcher_class_attribute(self):
+        environment = Environment(user_classes=[MyUserWithSameName1])
+
+        self.assertEqual(environment.dispatcher_class, UsersDispatcher)
+
+        class MyUsersDispatcher(UsersDispatcher):
+            pass
+
+        environment = Environment(user_classes=[MyUserWithSameName1], dispatcher_class=MyUsersDispatcher)
+
+        self.assertEqual(environment.dispatcher_class, MyUsersDispatcher)
+
     def test_update_user_class(self):
         class MyUser1(User):
             @task

locust/test/test_runners.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import locust
 from locust import (
     LoadTestShape,
@@ -6,6 +8,7 @@ from locust import (
     runners,
 )
 from locust.argument_parser import parse_options
+from locust.dispatch import UsersDispatcher
 from locust.env import Environment
 from locust.exception import RPCError, RPCReceiveError, StopUser
 from locust.main import create_environment
@@ -843,6 +846,42 @@ class TestLocustRunner(LocustRunnerTestCase):
         local_runner.stop()
         web_ui.stop()
 
+    def test_custom_dispatcher_class(self):
+        @mock.create_autospec
+        class MyUsersDispatcher(UsersDispatcher):
+            def __init__(self, worker_nodes: list[WorkerNode], user_classes: list[type[User]]) -> None:
+                super().__init__(worker_nodes, user_classes)
+
+            def add_worker(self, worker_node: WorkerNode) -> None:
+                pass
+
+            def remove_worker(self, worker_node: WorkerNode) -> None:
+                pass
+
+            @property
+            def dispatch_in_progress(self) -> bool:
+                return True
+
+            def new_dispatch(self, target_user_count: int, spawn_rate: float, user_classes: list | None = None) -> None:
+                return super().new_dispatch(target_user_count, spawn_rate, user_classes)
+
+        class MyUser1(User):
+            wait_time = constant(0)
+            fixed_count = 1
+
+            @task
+            def my_task(self):
+                pass
+
+        environment = Environment(user_classes=[MyUser1], dispatcher_class=MyUsersDispatcher)
+        runner = LocalRunner(environment)
+
+        runner.start(user_count=1, spawn_rate=1)
+        sleep(1)
+        runner._users_dispatcher.new_dispatch.assert_called_with(1, 1, None)
+        runner._users_dispatcher.__iter__.assert_called_once()
+        runner.stop()
+
 
 class TestMasterWorkerRunners(LocustTestCase):
     def test_distributed_integration_run(self):

locust/test/util.py

@@ -3,6 +3,7 @@ import functools
 import gc
 import os
 import socket
+import warnings
 from contextlib import contextmanager
 from tempfile import NamedTemporaryFile
 
@@ -78,7 +79,9 @@ def create_tls_cert(hostname):
 def clear_all_functools_lru_cache() -> None:
     # Clear all `functools.lru_cache` to ensure that no state are persisted from one test to another.
     # Taken from https://stackoverflow.com/a/50699209.
-    gc.collect()
+    with warnings.catch_warnings():
+        warnings.simplefilter(action="ignore", category=ResourceWarning)
+        gc.collect()
     wrappers = [a for a in gc.get_objects() if isinstance(a, functools._lru_cache_wrapper)]
     assert len(wrappers) > 0
     for wrapper in wrappers:

locust/user/task.py

@@ -308,7 +308,7 @@ class TaskSet(metaclass=TaskSetMeta):
         """Parent TaskSet instance of this TaskSet (or :py:class:`User <locust.User>` if this is not a nested TaskSet)"""
         return self._parent
 
-    def on_start(self):
+    def on_start(self) -> None:
         """
         Called when a User starts executing this TaskSet
         """

locust/user/users.py

@@ -128,7 +128,7 @@ class User(metaclass=UserMeta):
         self._taskset_instance: TaskSet = None
         self._cp_last_run = time.time()  # used by constant_pacing wait_time
 
-    def on_start(self):
+    def on_start(self) -> None:
         """
         Called when a User starts running.
         """