locust 2.23.2.dev7__tar.gz → 2.23.2.dev18__tar.gz

{locust-2.23.2.dev7 → locust-2.23.2.dev18}/.gitignore

@@ -28,3 +28,4 @@ __pycache__
 .env
 yarn-error.log
 .venv
+.DS_Store

{locust-2.23.2.dev7/locust.egg-info → locust-2.23.2.dev18}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: locust
-Version: 2.23.2.dev7
+Version: 2.23.2.dev18
 Summary: Developer friendly load testing framework
 License: MIT
 Project-URL: Homepage, https://locust.io/

{locust-2.23.2.dev7 → locust-2.23.2.dev18}/docs/configuration.rst

@@ -178,6 +178,7 @@ It is even possible to add custom arguments that you wish to be configurable for
 to your user:
 
 .. code-block:: python
+
     class Example(HttpUser):
         @task
         def example_task(self):

{locust-2.23.2.dev7 → locust-2.23.2.dev18}/docs/running-cloud-integration.rst

@@ -4,6 +4,10 @@
 Running Locust distributed with Terraform/AWS
 =============================================
 
+.. note::
+
+    This is broken at the moment. See https://github.com/locustio/locust/issues/2357 for a possible workaround.
+
 Here's one way to provision Locust using Infrastructure as Code/IaC.
 
 For this we use a `Terraform <https://www.terraform.io/>`_ module that provisions 1 leader and "n" worker nodes. This implementation uses AWS and EC2, but you may be able to modify it for other cloud providers.

{locust-2.23.2.dev7 → locust-2.23.2.dev18}/examples/terraform/aws/README.md

@@ -1,3 +1,4 @@
+## This doesnt seem to work right now for some people. See https://github.com/locustio/locust/issues/2357
 
 ## 1. AWS Authentication
 

{locust-2.23.2.dev7 → locust-2.23.2.dev18}/examples/web_ui_auth.py

@@ -6,7 +6,7 @@ according to your projects specifications.
 
 For more information, see https://docs.locust.io/en/stable/extending-locust.html#authentication
 """
-from locust import HttpUser, events
+from locust import HttpUser, events, task
 
 import json
 import os
@@ -16,7 +16,9 @@ from flask_login import UserMixin, login_user
 
 
 class LocustHttpUser(HttpUser):
-    pass
+    @task
+    def example(self):
+        self.client.get("/")
 
 
 class AuthUser(UserMixin):

{locust-2.23.2.dev7 → locust-2.23.2.dev18}/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-2.23.2.dev7 → locust-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-2.23.2.dev7 → locust-2.23.2.dev18}/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-2.23.2.dev7 → locust-2.23.2.dev18}/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-2.23.2.dev7 → locust-2.23.2.dev18}/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-2.23.2.dev7 → locust-2.23.2.dev18}/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-2.23.2.dev7 → locust-2.23.2.dev18}/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-2.23.2.dev7 → locust-2.23.2.dev18}/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-2.23.2.dev7 → locust-2.23.2.dev18}/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.
         """