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

locust/user/inspectuser.py

@@ -1,7 +1,8 @@
-from collections import defaultdict
+from __future__ import annotations
+
 import inspect
+from collections import defaultdict
 from json import dumps
-from typing import List, Type, Dict
 
 from .task import TaskSet
 from .users import User
@@ -49,11 +50,11 @@ def _print_task_ratio(x, level=0):
             _print_task_ratio(v["tasks"], level + 1)
 
 
-def get_ratio(user_classes: List[Type[User]], user_spawned: Dict[str, int], total: bool) -> Dict[str, Dict[str, float]]:
+def get_ratio(user_classes: list[type[User]], user_spawned: dict[str, int], total: bool) -> dict[str, dict[str, float]]:
     user_count = sum(user_spawned.values()) or 1
-    ratio_percent: Dict[Type[User], float] = {u: user_spawned.get(u.__name__, 0) / user_count for u in user_classes}
+    ratio_percent: dict[type[User], float] = {u: user_spawned.get(u.__name__, 0) / user_count for u in user_classes}
 
-    task_dict: Dict[str, Dict[str, float]] = {}
+    task_dict: dict[str, dict[str, float]] = {}
     for u, r in ratio_percent.items():
         d = {"ratio": r}
         d["tasks"] = _get_task_ratio(u.tasks, total, r)

locust/user/sequential_taskset.py

@@ -1,5 +1,7 @@
-import logging
 from locust.exception import LocustError
+
+import logging
+
 from .task import TaskSet, TaskSetMeta
 
 

locust/user/task.py

@@ -1,4 +1,7 @@
 from __future__ import annotations
+
+from locust.exception import InterruptTaskSet, MissingWaitTimeError, RescheduleTask, RescheduleTaskImmediately, StopUser
+
 import logging
 import random
 import traceback
@@ -6,22 +9,17 @@ from time import time
 from typing import (
     TYPE_CHECKING,
     Callable,
-    List,
-    TypeVar,
-    Optional,
-    Type,
-    overload,
-    Dict,
-    Set,
     Protocol,
+    Type,
+    TypeVar,
     final,
+    overload,
     runtime_checkable,
 )
+
 import gevent
 from gevent import GreenletExit
 
-from locust.exception import InterruptTaskSet, RescheduleTask, RescheduleTaskImmediately, StopUser, MissingWaitTimeError
-
 if TYPE_CHECKING:
     from locust import User
 
@@ -34,7 +32,7 @@ LOCUST_STATE_RUNNING, LOCUST_STATE_WAITING, LOCUST_STATE_STOPPING = ["running",
 
 @runtime_checkable
 class TaskHolder(Protocol[TaskT]):
-    tasks: List[TaskT]
+    tasks: list[TaskT]
 
 
 @overload
@@ -170,10 +168,10 @@ def get_tasks_from_base_classes(bases, class_dict):
 
 
 def filter_tasks_by_tags(
-    task_holder: Type[TaskHolder],
-    tags: Optional[Set[str]] = None,
-    exclude_tags: Optional[Set[str]] = None,
-    checked: Optional[Dict[TaskT, bool]] = None,
+    task_holder: type[TaskHolder],
+    tags: set[str] | None = None,
+    exclude_tags: set[str] | None = None,
+    checked: dict[TaskT, bool] | None = None,
 ):
     """
     Function used by Environment to recursively remove any tasks/TaskSets from a TaskSet/User that
@@ -238,7 +236,7 @@ class TaskSet(metaclass=TaskSetMeta):
     will then continue in the first TaskSet).
     """
 
-    tasks: List[TaskSet | Callable] = []
+    tasks: list[TaskSet | Callable] = []
     """
     Collection of python callables and/or TaskSet classes that the User(s) will run.
 
@@ -253,7 +251,7 @@ class TaskSet(metaclass=TaskSetMeta):
             tasks = {ThreadPage:15, write_post:1}
     """
 
-    min_wait: Optional[float] = None
+    min_wait: float | None = None
     """
     Deprecated: Use wait_time instead.
     Minimum waiting time between the execution of user tasks. Can be used to override
@@ -261,7 +259,7 @@ class TaskSet(metaclass=TaskSetMeta):
     TaskSet.
     """
 
-    max_wait: Optional[float] = None
+    max_wait: float | None = None
     """
     Deprecated: Use wait_time instead.
     Maximum waiting time between the execution of user tasks. Can be used to override
@@ -277,11 +275,11 @@ class TaskSet(metaclass=TaskSetMeta):
     if not set on the TaskSet.
     """
 
-    _user: "User"
-    _parent: "User"
+    _user: User
+    _parent: User
 
-    def __init__(self, parent: "User") -> None:
-        self._task_queue: List[Callable] = []
+    def __init__(self, parent: User) -> None:
+        self._task_queue: list[Callable] = []
         self._time_start = time()
 
         if isinstance(parent, TaskSet):
@@ -298,9 +296,10 @@ class TaskSet(metaclass=TaskSetMeta):
             self.max_wait = self.user.max_wait
         if not self.wait_function:
             self.wait_function = self.user.wait_function
+        self._cp_last_run = time()  # used by constant_pacing wait_time
 
     @property
-    def user(self) -> "User":
+    def user(self) -> User:
         """:py:class:`User <locust.User>` instance that this TaskSet was created by"""
         return self._user
 
@@ -426,11 +425,7 @@ class TaskSet(metaclass=TaskSetMeta):
             return random.randint(self.min_wait, self.max_wait) / 1000.0
         else:
             raise MissingWaitTimeError(
-                "You must define a wait_time method on either the %s or %s class"
-                % (
-                    type(self.user).__name__,
-                    type(self).__name__,
-                )
+                "You must define a wait_time method on either the {type(self.user).__name__} or {type(self).__name__} class"
             )
 
     def wait(self):

locust/user/users.py

@@ -1,12 +1,13 @@
 from __future__ import annotations
-from typing import Callable, Dict, List, Optional, final
 
+import logging
 import time
+import traceback
+from typing import Callable, final
+
 from gevent import GreenletExit, greenlet
 from gevent.pool import Group
 from urllib3 import PoolManager
-import logging
-import traceback
 
 logger = logging.getLogger(__name__)
 from locust.clients import HttpSession
@@ -59,7 +60,7 @@ class User(metaclass=UserMeta):
     :py:class:`HttpUser <locust.HttpUser>` class.
     """
 
-    host: Optional[str] = None
+    host: str | None = None
     """Base hostname to swarm. i.e: http://127.0.0.1:1234"""
 
     min_wait = None
@@ -89,7 +90,7 @@ class User(metaclass=UserMeta):
     Method that returns the time between the execution of locust tasks in milliseconds
     """
 
-    tasks: List[TaskSet | Callable] = []
+    tasks: list[TaskSet | Callable] = []
     """
     Collection of python callables and/or TaskSet classes that the Locust user(s) will run.
 
@@ -216,13 +217,22 @@ class User(metaclass=UserMeta):
     def greenlet(self):
         return self._greenlet
 
-    def context(self) -> Dict:
+    def context(self) -> dict:
         """
         Adds the returned value (a dict) to the context for :ref:`request event <request_context>`.
         Override this in your User class to customize the context.
         """
         return {}
 
+    @classmethod
+    def json(cls):
+        return {
+            "host": cls.host,
+            "weight": cls.weight,
+            "fixed_count": cls.fixed_count,
+            "tasks": [task.__name__ for task in cls.tasks],
+        }
+
     @classmethod
     def fullname(cls) -> str:
         """Fully qualified name of the user class, e.g. my_package.my_module.MyUserClass"""
@@ -244,7 +254,7 @@ class HttpUser(User):
     abstract = True
     """If abstract is True, the class is meant to be subclassed, and users will not choose this locust during a test"""
 
-    pool_manager: Optional[PoolManager] = None
+    pool_manager: PoolManager | None = None
     """Connection pool manager to use. If not given, a new manager is created per single user."""
 
     def __init__(self, *args, **kwargs):

locust/util/deprecation.py

@@ -1,6 +1,5 @@
 import warnings
 
-
 # Show deprecation warnings
 warnings.filterwarnings("always", category=DeprecationWarning, module="locust")
 

locust/util/exception_handler.py

@@ -1,5 +1,5 @@
-import time
 import logging
+import time
 
 logger = logging.getLogger(__name__)
 

locust/util/load_locustfile.py

@@ -1,8 +1,10 @@
+from __future__ import annotations
+
 import importlib
 import inspect
 import os
 import sys
-from typing import Dict, List, Optional, Tuple
+
 from ..shape import LoadTestShape
 from ..user import User
 
@@ -21,7 +23,7 @@ def is_shape_class(item):
     return bool(inspect.isclass(item) and issubclass(item, LoadTestShape) and not getattr(item, "abstract", True))
 
 
-def load_locustfile(path) -> Tuple[Optional[str], Dict[str, User], List[LoadTestShape]]:
+def load_locustfile(path) -> tuple[str | None, dict[str, User], list[LoadTestShape]]:
     """
     Import given locustfile path and return (docstring, callables).
 

locust/web.py

@@ -1,32 +1,44 @@
 from __future__ import annotations
+
 import csv
+import json
 import logging
 import os.path
 from functools import wraps
-
 from html import escape
 from io import StringIO
-from json import dumps
 from itertools import chain
+from json import dumps
 from time import time
-from typing import TYPE_CHECKING, Optional, Any, Dict, List
+from typing import TYPE_CHECKING, Any
 
 import gevent
-from flask import Flask, make_response, jsonify, render_template, request, send_file, Response, send_from_directory
-from flask_basicauth import BasicAuth
+from flask import (
+    Flask,
+    Response,
+    jsonify,
+    make_response,
+    redirect,
+    render_template,
+    request,
+    send_file,
+    send_from_directory,
+    url_for,
+)
+from flask_cors import CORS
+from flask_login import LoginManager, login_required
 from gevent import pywsgi
 
-from .exception import AuthCredentialsError
-from .runners import MasterRunner, STATE_RUNNING, STATE_MISSING
+from . import __version__ as version
+from . import argument_parser
+from . import stats as stats_module
+from .html import get_html_report
 from .log import greenlet_exception_logger
-from .stats import StatsCSVFileWriter, StatsErrorDict, sort_stats
-from . import stats as stats_module, __version__ as version, argument_parser
-from .stats import StatsCSV
+from .runners import STATE_MISSING, STATE_RUNNING, MasterRunner
+from .stats import StatsCSV, StatsCSVFileWriter, StatsErrorDict, sort_stats
 from .user.inspectuser import get_ratio
 from .util.cache import memoize
 from .util.timespan import parse_timespan
-from .html import get_html_report
-from flask_cors import CORS
 
 if TYPE_CHECKING:
     from .env import Environment
@@ -45,7 +57,7 @@ class WebUI:
     in :attr:`environment.stats <locust.env.Environment.stats>`
     """
 
-    app: Optional[Flask] = None
+    app: Flask | None = None
     """
     Reference to the :class:`flask.Flask` app. Can be used to add additional web routes and customize
     the Flask app in other various ways. Example::
@@ -57,27 +69,30 @@ class WebUI:
             return "your IP is: %s" % request.remote_addr
     """
 
-    greenlet: Optional[gevent.Greenlet] = None
+    greenlet: gevent.Greenlet | None = None
     """
     Greenlet of the running web server
     """
 
-    server: Optional[pywsgi.WSGIServer] = None
+    server: pywsgi.WSGIServer | None = None
     """Reference to the :class:`pyqsgi.WSGIServer` instance"""
 
-    template_args: Dict[str, Any]
+    template_args: dict[str, Any]
     """Arguments used to render index.html for the web UI. Must be used with custom templates
     extending index.html."""
 
+    auth_args: dict[str, Any]
+    """Arguments used to render auth.html for the web UI auth page. Must be used when configuring auth"""
+
     def __init__(
         self,
-        environment: "Environment",
+        environment: Environment,
         host: str,
         port: int,
-        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,
@@ -89,8 +104,7 @@ class WebUI:
         environment: Reference to the current Locust Environment
         host: Host/interface that the web server should accept connections to
         port: Port that the web server should listen to
-        auth_credentials:  If provided, it will enable basic auth with all the routes protected by default.
-                           Should be supplied in the format: "user:pass".
+        web_login:  Enables a login page for the modern UI
         tls_cert: A path to a TLS certificate
         tls_key: A path to a TLS private key
         delayed_start: Whether or not to delay starting web UI until `start()` is called. Delaying web UI start
@@ -105,6 +119,7 @@ class WebUI:
         self.tls_key = tls_key
         self.userclass_picker_is_active = userclass_picker_is_active
         self.modern_ui = modern_ui
+        self.web_login = web_login
         app = Flask(__name__)
         CORS(app)
         self.app = app
@@ -113,24 +128,16 @@ class WebUI:
         root_path = os.path.dirname(os.path.abspath(__file__))
         app.root_path = root_path
         self.webui_build_path = os.path.join(root_path, "webui", "dist")
-        self.app.config["BASIC_AUTH_ENABLED"] = False
-        self.auth: Optional[BasicAuth] = None
-        self.greenlet: Optional[gevent.Greenlet] = None
-        self._swarm_greenlet: Optional[gevent.Greenlet] = None
+        self.greenlet: gevent.Greenlet | None = None
+        self._swarm_greenlet: gevent.Greenlet | None = None
         self.template_args = {}
+        self.auth_args = {}
+
+        if self.web_login:
+            self.login_manager = LoginManager()
+            self.login_manager.init_app(app)
+            self.login_manager.login_view = "login"
 
-        if auth_credentials is not None:
-            credentials = auth_credentials.split(":")
-            if len(credentials) == 2:
-                self.app.config["BASIC_AUTH_USERNAME"] = credentials[0]
-                self.app.config["BASIC_AUTH_PASSWORD"] = credentials[1]
-                self.app.config["BASIC_AUTH_ENABLED"] = True
-                self.auth = BasicAuth()
-                self.auth.init_app(self.app)
-            else:
-                raise AuthCredentialsError(
-                    "Invalid auth_credentials. It should be a string in the following format: 'user:pass'"
-                )
         if environment.runner:
             self.update_template_args()
         if not delayed_start:
@@ -365,8 +372,8 @@ class WebUI:
         @self.auth_required_if_enabled
         @memoize(timeout=DEFAULT_CACHE_TIME, dynamic_timeout=True)
         def request_stats() -> Response:
-            stats: List[Dict[str, Any]] = []
-            errors: List[StatsErrorDict] = []
+            stats: list[dict[str, Any]] = []
+            errors: list[StatsErrorDict] = []
 
             if environment.runner is None:
                 report = {
@@ -476,9 +483,9 @@ class WebUI:
 
         @app.route("/tasks")
         @self.auth_required_if_enabled
-        def tasks() -> Dict[str, Dict[str, Dict[str, float]]]:
+        def tasks() -> dict[str, dict[str, dict[str, float]]]:
             runner = self.environment.runner
-            user_spawned: Dict[str, int]
+            user_spawned: dict[str, int]
             if runner is None:
                 user_spawned = {}
             else:
@@ -508,6 +515,30 @@ class WebUI:
 
             return jsonify({"logs": logs})
 
+        @app.route("/login")
+        def login():
+            if not self.web_login:
+                return redirect(url_for("index"))
+
+            if self.modern_ui:
+                self.set_static_modern_ui()
+
+                return render_template(
+                    "auth.html",
+                    auth_args=self.auth_args,
+                )
+            else:
+                return "Web Auth is only available on the modern web ui. Enable it with the --modern-ui flag"
+
+        @app.route("/user", methods=["POST"])
+        def update_user():
+            assert request.method == "POST"
+
+            user_settings = json.loads(request.data)
+            self.environment.update_user_class(user_settings)
+
+            return {}, 201
+
     def start(self):
         self.greenlet = gevent.spawn(self.start_server)
         self.greenlet.link_exception(greenlet_exception_handler)
@@ -529,8 +560,8 @@ class WebUI:
 
     def auth_required_if_enabled(self, view_func):
         """
-        Decorator that can be used on custom route methods that will turn on Basic Auth
-        authentication if the ``--web-auth`` flag is used. Example::
+        Decorator that can be used on custom route methods that will turn on Flask Login
+        authentication if the ``--web-login`` flag is used. Example::
 
             @web_ui.app.route("/my_custom_route")
             @web_ui.auth_required_if_enabled
@@ -540,11 +571,11 @@ class WebUI:
 
         @wraps(view_func)
         def wrapper(*args, **kwargs):
-            if self.app.config["BASIC_AUTH_ENABLED"]:
-                if self.auth.authenticate():
-                    return view_func(*args, **kwargs)
-                else:
-                    return self.auth.challenge()
+            if self.web_login:
+                try:
+                    return login_required(view_func)(*args, **kwargs)
+                except Exception as e:
+                    return f"Locust auth exception: {e} See https://docs.locust.io/en/stable/extending-locust.html#authentication for configuring authentication."
             else:
                 return view_func(*args, **kwargs)
 
@@ -582,17 +613,32 @@ class WebUI:
         stats = self.environment.runner.stats
         extra_options = argument_parser.ui_extra_args_dict()
 
-        available_user_classes = (
-            None if not self.environment.available_user_classes else sorted(self.environment.available_user_classes)
-        )
+        available_user_classes = None
+        users = None
+        if self.environment.available_user_classes:
+            available_user_classes = sorted(self.environment.available_user_classes)
+            users = {
+                user_class_name: user_class.json()
+                for (user_class_name, user_class) in self.environment.available_user_classes.items()
+            }
 
         available_shape_classes = ["Default"]
         if self.environment.available_shape_classes:
             available_shape_classes += sorted(self.environment.available_shape_classes.keys())
 
+        available_user_tasks = (
+            {
+                user_class_name: [task.__name__ for task in user_class]
+                for (user_class_name, user_class) in self.environment.available_user_tasks.items()
+            }
+            if self.environment.available_user_tasks
+            else None
+        )
+
         if self.modern_ui:
             percentiles = {
                 "percentiles_to_chart": stats_module.MODERN_UI_PERCENTILES_TO_CHART,
+                "percentiles_to_statistics": stats_module.PERCENTILES_TO_STATISTICS,
             }
         else:
             percentiles = {
@@ -623,12 +669,15 @@ class WebUI:
             "show_userclass_picker": self.userclass_picker_is_active,
             "available_user_classes": available_user_classes,
             "available_shape_classes": available_shape_classes,
+            "available_user_tasks": available_user_tasks,
+            "users": users,
             **percentiles,
         }
 
     def _update_shape_class(self, shape_class_name):
         if shape_class_name:
             shape_class = self.environment.available_shape_classes[shape_class_name]
+            shape_class.runner = self.environment.runner
         else:
             shape_class = None
 

locust/webui/dist/assets/auth-5e21717c.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"auth-5e21717c.js","sources":[],"sourcesContent":[],"names":[],"mappings":""}