locust 2.20.1.dev26__py3-none-any.whl → 2.20.2__py3-none-any.whl

locust/env.py

@@ -1,24 +1,18 @@
+from __future__ import annotations
+
 from operator import methodcaller
-from typing import (
-    Callable,
-    Dict,
-    List,
-    Type,
-    TypeVar,
-    Optional,
-)
+from typing import Callable, TypeVar
 
 from configargparse import Namespace
 
 from .event import Events
 from .exception import RunnerAlreadyExistsError
+from .runners import LocalRunner, MasterRunner, Runner, WorkerRunner
+from .shape import LoadTestShape
 from .stats import RequestStats, StatsCSV
-from .runners import Runner, LocalRunner, MasterRunner, WorkerRunner
-from .web import WebUI
 from .user import User
-from .user.task import filter_tasks_by_tags, TaskSet, TaskHolder
-from .shape import LoadTestShape
-
+from .user.task import TaskHolder, TaskSet, filter_tasks_by_tags
+from .web import WebUI
 
 RunnerType = TypeVar("RunnerType", bound=Runner)
 
@@ -27,27 +21,28 @@ class Environment:
     def __init__(
         self,
         *,
-        user_classes: Optional[List[Type[User]]] = None,
-        shape_class: Optional[LoadTestShape] = None,
-        tags: Optional[List[str]] = None,
-        locustfile: Optional[str] = None,
-        exclude_tags: Optional[List[str]] = None,
-        events: Optional[Events] = None,
-        host: Optional[str] = None,
+        user_classes: list[type[User]] | None = None,
+        shape_class: LoadTestShape | None = None,
+        tags: list[str] | None = None,
+        locustfile: str | None = None,
+        exclude_tags: list[str] | None = None,
+        events: Events | None = None,
+        host: str | None = None,
         reset_stats=False,
-        stop_timeout: Optional[float] = None,
+        stop_timeout: float | None = None,
         catch_exceptions=True,
-        parsed_options: Optional[Namespace] = None,
-        available_user_classes: Optional[Dict[str, User]] = None,
-        available_shape_classes: Optional[Dict[str, LoadTestShape]] = None,
+        parsed_options: Namespace | None = None,
+        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,
     ):
-        self.runner: Optional[Runner] = None
+        self.runner: Runner | None = None
         """Reference to the :class:`Runner <locust.runners.Runner>` instance"""
 
-        self.web_ui: Optional[WebUI] = None
+        self.web_ui: WebUI | None = None
         """Reference to the WebUI instance"""
 
-        self.process_exit_code: Optional[int] = None
+        self.process_exit_code: int | None = None
         """
         If set it'll be the exit code of the Locust process
         """
@@ -63,7 +58,7 @@ class Environment:
 
         self.locustfile = locustfile
         """Filename (not path) of locustfile"""
-        self.user_classes: List[Type[User]] = user_classes or []
+        self.user_classes: list[type[User]] = user_classes or []
         """User classes that the runner will run"""
         self.shape_class = shape_class
         """A shape class to control the shape of the load test"""
@@ -98,6 +93,8 @@ class Environment:
         """List of the available User Classes to pick from in the UserClass Picker"""
         self.available_shape_classes = available_shape_classes
         """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._remove_user_classes_with_weight_zero()
         self._validate_user_class_name_uniqueness()
@@ -105,7 +102,7 @@ class Environment:
 
     def _create_runner(
         self,
-        runner_class: Type[RunnerType],
+        runner_class: type[RunnerType],
         *args,
         **kwargs,
     ) -> RunnerType:
@@ -159,10 +156,10 @@ class Environment:
         self,
         host="",
         port=8089,
-        auth_credentials: Optional[str] = None,
-        tls_cert: Optional[str] = None,
-        tls_key: Optional[str] = None,
-        stats_csv_writer: Optional[StatsCSV] = None,
+        web_login: bool = False,
+        tls_cert: str | None = None,
+        tls_key: str | None = None,
+        stats_csv_writer: StatsCSV | None = None,
         delayed_start=False,
         userclass_picker_is_active=False,
         modern_ui=False,
@@ -173,7 +170,7 @@ class Environment:
         :param host: Host/interface that the web server should accept connections to. Defaults to ""
                      which means all interfaces
         :param port: Port that the web server should listen to
-        :param auth_credentials: If provided (in format "username:password") basic auth will be enabled
+        :param web_login: If provided, an authentication page will protect the app
         :param tls_cert: An optional path (str) to a TLS cert. If this is provided the web UI will be
                          served over HTTPS
         :param tls_key: An optional path (str) to a TLS private key. If this is provided the web UI will be
@@ -186,7 +183,7 @@ class Environment:
             self,
             host,
             port,
-            auth_credentials=auth_credentials,
+            web_login=web_login,
             tls_cert=tls_cert,
             tls_key=tls_key,
             stats_csv_writer=stats_csv_writer,
@@ -196,6 +193,17 @@ class Environment:
         )
         return self.web_ui
 
+    def update_user_class(self, user_settings):
+        user_class_name = user_settings.get("user_class_name")
+        user_class = self.available_user_classes[user_class_name]
+        user_tasks = self.available_user_tasks[user_class_name]
+
+        for key, value in user_settings.items():
+            if key not in ["user_class_name", "tasks"]:
+                setattr(user_class, key, value)
+            if key == "tasks":
+                user_class.tasks = [task for task in user_tasks if task.__name__ in value]
+
     def _filter_tasks_by_tags(self) -> None:
         """
         Filter the tasks on all the user_classes recursively, according to the tags and
@@ -244,7 +252,7 @@ class Environment:
         """
         for u in self.user_classes:
             u.weight = 1
-            user_tasks: List[TaskSet | Callable] = []
+            user_tasks: list[TaskSet | Callable] = []
             tasks_frontier = u.tasks
             while len(tasks_frontier) != 0:
                 t = tasks_frontier.pop()
@@ -274,5 +282,5 @@ class Environment:
             )
 
     @property
-    def user_classes_by_name(self) -> Dict[str, Type[User]]:
+    def user_classes_by_name(self) -> dict[str, type[User]]:
         return {u.__name__: u for u in self.user_classes}

locust/event.py

@@ -1,10 +1,13 @@
+from __future__ import annotations
+
 import logging
-from . import log
+import time
 import traceback
 from contextlib import contextmanager
-from typing import Generator, Any, Dict
-import time
-from .exception import StopUser, RescheduleTask, RescheduleTaskImmediately, InterruptTaskSet
+from typing import Any, Generator
+
+from . import log
+from .exception import InterruptTaskSet, RescheduleTask, RescheduleTaskImmediately, StopUser
 
 
 class EventHook:
@@ -52,7 +55,7 @@ class EventHook:
     @contextmanager
     def measure(
         self, request_type: str, name: str, response_length: int = 0, context=None
-    ) -> Generator[Dict[str, Any], None, None]:
+    ) -> Generator[dict[str, Any], None, None]:
         """Convenience method for firing the event with automatically calculated response time and automatically marking the request as failed if an exception is raised (this is really only useful for the *request* event)
 
         Example usage (in a task):
@@ -230,11 +233,11 @@ class Events:
     """
 
     def __init__(self):
-        # For backwarde compatibility use also values of class attributes
+        # For backward compatibility use also values of class attributes
         for name, value in vars(type(self)).items():
-            if value == EventHook:
-                setattr(self, name, value())
+            if value == "EventHook":
+                setattr(self, name, EventHook())
 
         for name, value in self.__annotations__.items():
-            if value == EventHook:
-                setattr(self, name, value())
+            if value == "EventHook":
+                setattr(self, name, EventHook())

locust/exception.py

@@ -74,14 +74,5 @@ class RPCReceiveError(Exception):
         self.addr = addr
 
 
-class AuthCredentialsError(ValueError):
-    """
-    Exception when the auth credentials provided
-    are not in the correct format
-    """
-
-    pass
-
-
 class RunnerAlreadyExistsError(Exception):
     pass

locust/html.py

@@ -1,16 +1,17 @@
-from jinja2 import Environment, FileSystemLoader
-import os
+import datetime
 import glob
+import os
 import pathlib
-import datetime
-from itertools import chain
-from .stats import sort_stats
-from . import stats as stats_module
-from .user.inspectuser import get_ratio
 from html import escape
+from itertools import chain
 from json import dumps
-from .runners import MasterRunner, STATE_STOPPED, STATE_STOPPING
 
+from jinja2 import Environment, FileSystemLoader
+
+from . import stats as stats_module
+from .runners import STATE_STOPPED, STATE_STOPPING, MasterRunner
+from .stats import sort_stats
+from .user.inspectuser import get_ratio
 
 PERCENTILES_FOR_HTML_REPORT = [0.50, 0.6, 0.7, 0.8, 0.9, 0.95, 0.99, 1.0]
 

locust/input_events.py

@@ -1,18 +1,20 @@
-from typing import Dict, Callable
+from __future__ import annotations
 
-import gevent
 import logging
 import os
 import sys
+from typing import Callable
+
+import gevent
 
 if os.name == "nt":
     from win32api import STD_INPUT_HANDLE
     from win32console import (
-        GetStdHandle,
-        KEY_EVENT,
         ENABLE_ECHO_INPUT,
         ENABLE_LINE_INPUT,
         ENABLE_PROCESSED_INPUT,
+        KEY_EVENT,
+        GetStdHandle,
     )
 else:
     import select
@@ -88,7 +90,7 @@ def get_poller():
         return UnixKeyPoller()
 
 
-def input_listener(key_to_func_map: Dict[str, Callable]):
+def input_listener(key_to_func_map: dict[str, Callable]):
     def input_listener_func():
         try:
             poller = get_poller()

locust/main.py

@@ -1,35 +1,40 @@
+from __future__ import annotations
+
+import locust
+
+import atexit
 import errno
+import gc
+import inspect
+import json
 import logging
 import os
 import signal
 import sys
 import time
-import atexit
-import inspect
+import traceback
+
 import gevent
-import locust
-from typing import Dict
-from . import log
+
+from . import log, stats
 from .argument_parser import parse_locustfile_option, parse_options
 from .env import Environment
-from .log import setup_logging, greenlet_exception_logger
-from . import stats
+from .html import get_html_report
+from .input_events import input_listener
+from .log import greenlet_exception_logger, setup_logging
 from .stats import (
+    StatsCSV,
+    StatsCSVFileWriter,
     print_error_report,
     print_percentile_stats,
     print_stats,
     print_stats_json,
-    stats_printer,
     stats_history,
+    stats_printer,
 )
-from .stats import StatsCSV, StatsCSVFileWriter
 from .user.inspectuser import print_task_ratio, print_task_ratio_json
-from .util.timespan import parse_timespan
-from .exception import AuthCredentialsError
-from .input_events import input_listener
-from .html import get_html_report
 from .util.load_locustfile import load_locustfile
-import traceback
+from .util.timespan import parse_timespan
 
 try:
     # import locust_plugins if it is installed, to allow it to register custom arguments etc
@@ -56,6 +61,7 @@ def create_environment(
     locustfile=None,
     available_user_classes=None,
     available_shape_classes=None,
+    available_user_tasks=None,
 ):
     """
     Create an Environment instance from options
@@ -70,6 +76,7 @@ def create_environment(
         parsed_options=options,
         available_user_classes=available_user_classes,
         available_shape_classes=available_shape_classes,
+        available_user_tasks=available_user_tasks,
     )
 
 
@@ -84,9 +91,10 @@ def main():
     locustfile = locustfiles[0] if locustfiles_length == 1 else None
 
     # Importing Locustfile(s) - setting available UserClasses and ShapeClasses to choose from in UI
-    user_classes: Dict[str, locust.User] = {}
+    user_classes: dict[str, locust.User] = {}
     available_user_classes = {}
     available_shape_classes = {}
+    available_user_tasks = {}
     shape_class = None
     for _locustfile in locustfiles:
         docstring, _user_classes, shape_classes = load_locustfile(_locustfile)
@@ -118,6 +126,7 @@ def main():
 
             user_classes[key] = value
             available_user_classes[key] = value
+            available_user_tasks[key] = value.tasks or None
 
     if len(stats.PERCENTILES_TO_CHART) != 2:
         logging.error("stats.PERCENTILES_TO_CHART parameter should be 2 parameters \n")
@@ -137,6 +146,14 @@ def main():
                 "stats.PERCENTILES_TO_CHART parameter need to be float and value between. 0 < percentile < 1 Eg 0.95\n"
             )
             sys.exit(1)
+
+    for percentile in stats.PERCENTILES_TO_STATISTICS:
+        if not is_valid_percentile(percentile):
+            logging.error(
+                "stats.PERCENTILES_TO_STATISTICS parameter need to be float and value between. 0 < percentile < 1 Eg 0.95\n"
+            )
+            sys.exit(1)
+
     # parse all command line options
     options = parse_options()
 
@@ -147,6 +164,12 @@ def main():
         sys.stderr.write("The --slave/--expect-slaves parameters have been renamed --worker/--expect-workers\n")
         sys.exit(1)
 
+    if options.web_auth:
+        sys.stderr.write(
+            "The --web-auth parameters has been replaced with --web-login. See https://docs.locust.io/en/stable/extending-locust.html#authentication for details\n"
+        )
+        sys.exit(1)
+
     if options.autoquit != -1 and not options.autostart:
         sys.stderr.write("--autoquit is only meaningful in combination with --autostart\n")
         sys.exit(1)
@@ -182,6 +205,9 @@ def main():
                 "--master cannot be combined with --processes. Remove --master, as it is implicit as long as --worker is not set.\n"
             )
             sys.exit(1)
+        # Optimize copy-on-write-behavior to save some memory (aprx 26MB -> 15MB rss) in child processes
+        gc.collect()  # avoid freezing garbage
+        gc.freeze()  # move all objects to perm gen so ref counts dont get updated
         for _ in range(options.processes):
             child_pid = gevent.fork()
             if child_pid:
@@ -337,8 +363,36 @@ See https://github.com/locustio/locust/wiki/Installation#increasing-maximum-numb
         locustfile=locustfile_path,
         available_user_classes=available_user_classes,
         available_shape_classes=available_shape_classes,
+        available_user_tasks=available_user_tasks,
     )
 
+    if options.config_users:
+        for json_user_config in options.config_users:
+            try:
+                if json_user_config.endswith(".json"):
+                    with open(json_user_config) as file:
+                        user_config = json.load(file)
+                else:
+                    user_config = json.loads(json_user_config)
+
+                def ensure_user_class_name(config):
+                    if "user_class_name" not in config:
+                        logger.error("The user config must specify a user_class_name")
+                        sys.exit(-1)
+
+                if isinstance(user_config, list):
+                    for config in user_config:
+                        ensure_user_class_name(config)
+
+                        environment.update_user_class(config)
+                else:
+                    ensure_user_class_name(user_config)
+
+                    environment.update_user_class(user_config)
+            except Exception as e:
+                logger.error(f"The --config-users arugment must be in valid JSON string or file: {e}")
+                sys.exit(-1)
+
     if (
         shape_class
         and not shape_class.use_common_options
@@ -412,39 +466,33 @@ See https://github.com/locustio/locust/wiki/Installation#increasing-maximum-numb
     if not options.headless and not options.worker:
         # spawn web greenlet
         protocol = "https" if options.tls_cert and options.tls_key else "http"
-        try:
-            if options.web_host == "*":
-                # special check for "*" so that we're consistent with --master-bind-host
-                web_host = ""
-            else:
-                web_host = options.web_host
-            if web_host:
-                logger.info(f"Starting web interface at {protocol}://{web_host}:{options.web_port}")
-            else:
-                if os.name == "nt":
-                    logger.info(
-                        f"Starting web interface at {protocol}://localhost:{options.web_port} (accepting connections from all network interfaces)"
-                    )
-                else:
-                    logger.info(f"Starting web interface at {protocol}://0.0.0.0:{options.web_port}")
-            if options.web_auth:
-                logging.warning(
-                    "BasicAuth support is deprecated, it will be removed in a future release, unless someone reimplements it in a more modern way! See https://github.com/locustio/locust/issues/2517"
+
+        if options.web_host == "*":
+            # special check for "*" so that we're consistent with --master-bind-host
+            web_host = ""
+        else:
+            web_host = options.web_host
+        if web_host:
+            logger.info(f"Starting web interface at {protocol}://{web_host}:{options.web_port}")
+        else:
+            if os.name == "nt":
+                logger.info(
+                    f"Starting web interface at {protocol}://localhost:{options.web_port} (accepting connections from all network interfaces)"
                 )
-            web_ui = environment.create_web_ui(
-                host=web_host,
-                port=options.web_port,
-                auth_credentials=options.web_auth,
-                tls_cert=options.tls_cert,
-                tls_key=options.tls_key,
-                stats_csv_writer=stats_csv_writer,
-                delayed_start=True,
-                userclass_picker_is_active=options.class_picker,
-                modern_ui=options.modern_ui,
-            )
-        except AuthCredentialsError:
-            logger.error("Credentials supplied with --web-auth should have the format: username:password")
-            sys.exit(1)
+            else:
+                logger.info(f"Starting web interface at {protocol}://0.0.0.0:{options.web_port}")
+
+        web_ui = environment.create_web_ui(
+            host=web_host,
+            port=options.web_port,
+            web_login=options.web_login,
+            tls_cert=options.tls_cert,
+            tls_key=options.tls_key,
+            stats_csv_writer=stats_csv_writer,
+            delayed_start=True,
+            userclass_picker_is_active=options.class_picker,
+            modern_ui=options.modern_ui,
+        )
     else:
         web_ui = None
 

locust/rpc/protocol.py

@@ -1,6 +1,8 @@
-import msgpack
+from __future__ import annotations
+
 import datetime
-from typing import Type
+
+import msgpack
 
 try:
     from bson import ObjectId  # type: ignore

locust/rpc/zmqrpc.py

@@ -1,9 +1,11 @@
-import zmq.green as zmq
-from .protocol import Message
+from locust.exception import RPCError, RPCReceiveError, RPCSendError
 from locust.util.exception_handler import retry
-from locust.exception import RPCError, RPCSendError, RPCReceiveError
-import zmq.error as zmqerr
+
 import msgpack.exceptions as msgerr
+import zmq.error as zmqerr
+import zmq.green as zmq
+
+from .protocol import Message
 
 
 class BaseSocket: