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

locust/test/test_interruptable_task.py

@@ -1,11 +1,11 @@
-from collections import defaultdict
-from unittest import TestCase
-from typing import DefaultDict
-
 from locust import SequentialTaskSet, User, constant, task
 from locust.env import Environment
 from locust.exception import StopUser
 
+from collections import defaultdict
+from typing import DefaultDict
+from unittest import TestCase
+
 
 class InterruptableTaskSet(SequentialTaskSet):
     counter: DefaultDict[str, int] = defaultdict(int)

locust/test/test_load_locustfile.py

@@ -1,12 +1,13 @@
-import os
-import textwrap
-
 from locust import main
 from locust.argument_parser import parse_options
 from locust.main import create_environment
-from locust.user import HttpUser, User, TaskSet
+from locust.user import HttpUser, TaskSet, User
 from locust.util.load_locustfile import is_user_class
-from .mock_locustfile import mock_locustfile, MOCK_LOCUSTFILE_CONTENT
+
+import os
+import textwrap
+
+from .mock_locustfile import MOCK_LOCUSTFILE_CONTENT, mock_locustfile
 from .testcases import LocustTestCase
 from .util import temporary_file
 

locust/test/test_locust_class.py

@@ -1,12 +1,18 @@
+from locust import HttpUser, TaskSet, User, constant, task
+from locust.env import Environment
+from locust.exception import (
+    CatchResponseError,
+    InterruptTaskSet,
+    RescheduleTask,
+    RescheduleTaskImmediately,
+    ResponseError,
+    StopUser,
+)
+
 import gevent
 from gevent import sleep
 from gevent.pool import Group
 
-from locust.exception import InterruptTaskSet, ResponseError
-from locust import HttpUser, User, TaskSet, task, constant
-from locust.env import Environment
-from locust.exception import CatchResponseError, RescheduleTask, RescheduleTaskImmediately, StopUser
-
 from .testcases import LocustTestCase, WebserverTestCase
 
 

locust/test/test_log.py

@@ -1,16 +1,17 @@
-from unittest import mock
+from locust import log
+from locust.log import greenlet_exception_logger
+
 import socket
 import subprocess
 import textwrap
 from logging import getLogger
+from unittest import mock
 
 import gevent
 
-from locust import log
-from locust.log import greenlet_exception_logger
+from . import changed_rlimit
 from .testcases import LocustTestCase
 from .util import temporary_file
-from . import changed_rlimit
 
 
 class TestGreenletExceptionLogger(LocustTestCase):

locust/test/test_main.py

@@ -2,23 +2,24 @@ from __future__ import annotations
 
 import json
 import os
-import sys
 import platform
-import unittest
-import socket
-import psutil
 import pty
 import signal
+import socket
 import subprocess
+import sys
 import textwrap
+import unittest
+from subprocess import DEVNULL, PIPE, STDOUT
 from tempfile import TemporaryDirectory
 from unittest import TestCase
-from subprocess import PIPE, STDOUT, DEVNULL
+
 import gevent
+import psutil
 import requests
 
-from .mock_locustfile import mock_locustfile, MOCK_LOCUSTFILE_CONTENT
-from .util import temporary_file, get_free_tcp_port, patch_env
+from .mock_locustfile import MOCK_LOCUSTFILE_CONTENT, mock_locustfile
+from .util import get_free_tcp_port, patch_env, temporary_file
 
 
 def is_port_in_use(port: int) -> bool:
@@ -233,6 +234,35 @@ class StandaloneIntegrationTests(ProcessIntegrationTest):
             stdout, stderr = proc.communicate()
             self.assertIn("Starting web interface at", stderr)
 
+    def test_percentiles_to_statistics(self):
+        port = get_free_tcp_port()
+        with temporary_file(
+            content=textwrap.dedent(
+                """
+                from locust import User, task, constant, events
+                from locust.stats import PERCENTILES_TO_STATISTICS
+                PERCENTILES_TO_STATISTICS = [0.9, 0.99]
+                class TestUser(User):
+                    wait_time = constant(3)
+                    @task
+                    def my_task(self):
+                        print("running my_task()")
+            """
+            )
+        ) as file_path:
+            proc = subprocess.Popen(
+                ["locust", "-f", file_path, "--web-port", str(port), "--autostart", "--modern-ui"],
+                stdout=PIPE,
+                stderr=PIPE,
+                text=True,
+            )
+            gevent.sleep(1)
+            response = requests.get(f"http://localhost:{port}/")
+            self.assertEqual(200, response.status_code)
+            proc.send_signal(signal.SIGTERM)
+            stdout, stderr = proc.communicate()
+            self.assertIn("Starting web interface at", stderr)
+
     def test_invalid_percentile_parameter(self):
         with temporary_file(
             content=textwrap.dedent(

locust/test/test_parser.py

@@ -1,19 +1,20 @@
-import unittest
-import os
-from tempfile import NamedTemporaryFile, TemporaryDirectory
-from random import randint
-from unittest import mock
-from io import StringIO
-
 import locust
 from locust.argument_parser import (
-    locustfile_is_directory,
-    parse_options,
+    find_locustfiles,
     get_parser,
+    locustfile_is_directory,
     parse_locustfile_option,
+    parse_options,
     ui_extra_args_dict,
-    find_locustfiles,
 )
+
+import os
+import unittest
+from io import StringIO
+from random import randint
+from tempfile import NamedTemporaryFile, TemporaryDirectory
+from unittest import mock
+
 from .mock_locustfile import mock_locustfile
 from .testcases import LocustTestCase
 
@@ -49,11 +50,6 @@ class TestParser(unittest.TestCase):
         self.assertEqual(options.locustfile, "locustfile_from_env")
         self.assertEqual(options.host, "host_from_args")  # overridden
 
-    def test_web_auth(self):
-        args = ["--web-auth", "hello:bye"]
-        opts = self.parser.parse_args(args)
-        self.assertEqual(opts.web_auth, "hello:bye")
-
 
 class TestArgumentParser(LocustTestCase):
     def setUp(self):

locust/test/test_runners.py

@@ -1,48 +1,49 @@
-import json
-import random
-import time
-import unittest
-from collections import defaultdict
-from operator import itemgetter
-
-import gevent
-from unittest import mock
-import requests
-from gevent import sleep
-from gevent.pool import Group
-from gevent.queue import Queue
-
 import locust
 from locust import (
     LoadTestShape,
+    __version__,
     constant,
     runners,
-    __version__,
 )
 from locust.argument_parser import parse_options
 from locust.env import Environment
-from locust.exception import RPCError, StopUser, RPCReceiveError
+from locust.exception import RPCError, RPCReceiveError, StopUser
 from locust.main import create_environment
 from locust.rpc import Message
 from locust.runners import (
-    LocalRunner,
     STATE_INIT,
-    STATE_SPAWNING,
-    STATE_RUNNING,
     STATE_MISSING,
-    STATE_STOPPING,
+    STATE_RUNNING,
+    STATE_SPAWNING,
     STATE_STOPPED,
+    STATE_STOPPING,
+    LocalRunner,
     WorkerNode,
     WorkerRunner,
 )
 from locust.stats import RequestStats
-from .testcases import LocustTestCase
 from locust.user import (
     TaskSet,
     User,
     task,
 )
+
+import json
+import random
+import time
+import unittest
+from collections import defaultdict
+from operator import itemgetter
+from unittest import mock
+
+import gevent
+import requests
+from gevent import sleep
+from gevent.pool import Group
+from gevent.queue import Queue
 from retry import retry  # type: ignore
+
+from .testcases import LocustTestCase
 from .util import patch_env
 
 NETWORK_BROKEN = "network broken"

locust/test/test_sequential_taskset.py

@@ -1,6 +1,7 @@
-from locust import User, task, constant
-from locust.user.sequential_taskset import SequentialTaskSet
+from locust import User, constant, task
 from locust.exception import RescheduleTask
+from locust.user.sequential_taskset import SequentialTaskSet
+
 from .testcases import LocustTestCase
 
 

locust/test/test_stats.py

@@ -1,33 +1,31 @@
-import csv
-import time
-import unittest
-import re
-import os
-import json
-
-import gevent
-from unittest import mock
 import locust
-from locust import HttpUser, TaskSet, task, User, constant, __version__
+from locust import HttpUser, TaskSet, User, __version__, constant, task
 from locust.env import Environment
 from locust.rpc.protocol import Message
 from locust.stats import (
+    PERCENTILES_TO_REPORT,
+    STATS_NAME_WIDTH,
+    STATS_TYPE_WIDTH,
     CachedResponseTimes,
     RequestStats,
+    StatsCSVFileWriter,
     StatsEntry,
     diff_response_time_dicts,
-    PERCENTILES_TO_REPORT,
-    STATS_NAME_WIDTH,
-    STATS_TYPE_WIDTH,
+    stats_history,
 )
-from locust.stats import StatsCSVFileWriter
-from locust.stats import stats_history
-from locust.test.testcases import LocustTestCase
+from locust.test.test_runners import mocked_rpc
+from locust.test.testcases import LocustTestCase, WebserverTestCase
 from locust.user.inspectuser import _get_task_ratio
 
-from locust.test.testcases import WebserverTestCase
-from locust.test.test_runners import mocked_rpc
+import csv
+import json
+import os
+import re
+import time
+import unittest
+from unittest import mock
 
+import gevent
 
 _TEST_CSV_STATS_INTERVAL_SEC = 0.2
 _TEST_CSV_STATS_INTERVAL_WAIT_SEC = _TEST_CSV_STATS_INTERVAL_SEC + 0.1

locust/test/test_tags.py

@@ -1,6 +1,7 @@
-from locust import TaskSet, User, task, tag
-from locust.user.task import filter_tasks_by_tags
+from locust import TaskSet, User, tag, task
 from locust.env import Environment
+from locust.user.task import filter_tasks_by_tags
+
 from .testcases import LocustTestCase
 
 

locust/test/test_taskratio.py

@@ -1,7 +1,7 @@
-import unittest
+from locust.user import TaskSet, User, task
+from locust.user.inspectuser import _get_task_ratio, get_ratio
 
-from locust.user import User, TaskSet, task
-from locust.user.inspectuser import get_ratio, _get_task_ratio
+import unittest
 
 
 class TestTaskRatio(unittest.TestCase):

locust/test/test_users.py

@@ -1,10 +1,10 @@
+from locust import HttpUser, User
+from locust.test.testcases import WebserverTestCase
+
 import unittest
 
 from urllib3 import PoolManager
 
-from locust import User, HttpUser
-from locust.test.testcases import WebserverTestCase
-
 
 class TestUserClass(unittest.TestCase):
     class MyClassScopedUser(User):

locust/test/test_util.py

@@ -1,6 +1,7 @@
-import unittest
-from locust.util.timespan import parse_timespan
 from locust.util.rounding import proper_round
+from locust.util.timespan import parse_timespan
+
+import unittest
 
 
 class TestParseTimespan(unittest.TestCase):

locust/test/test_wait_time.py

@@ -1,9 +1,9 @@
+from locust import TaskSet, User, between, constant, constant_throughput
+from locust.exception import MissingWaitTimeError
+
 import random
 import time
 
-from locust import User, TaskSet, between, constant, constant_throughput
-from locust.exception import MissingWaitTimeError
-
 from .testcases import LocustTestCase
 
 

locust/test/test_web.py

@@ -1,32 +1,35 @@
+from __future__ import annotations
+
+import locust
+from locust import LoadTestShape, constant, stats
+from locust.argument_parser import get_parser, parse_options
+from locust.env import Environment
+from locust.log import LogReader
+from locust.runners import Runner
+from locust.stats import StatsCSVFileWriter
+from locust.user import User, task
+from locust.web import WebUI
+
 import copy
 import csv
 import json
+import logging
 import os
 import re
 import textwrap
 import traceback
-import logging
 from io import StringIO
 from tempfile import NamedTemporaryFile, TemporaryDirectory
 
 import gevent
 import requests
+from flask_login import UserMixin
 from pyquery import PyQuery as pq
-import locust
-from locust import constant, LoadTestShape
-from locust.argument_parser import get_parser, parse_options
-from locust.user import User, task
-from locust.env import Environment
-from locust.runners import Runner
-from locust import stats
-from locust.stats import StatsCSVFileWriter
-from locust.web import WebUI
-from locust.log import LogReader
 
+from ..util.load_locustfile import load_locustfile
 from .mock_locustfile import mock_locustfile
 from .testcases import LocustTestCase
 from .util import create_tls_cert
-from ..util.load_locustfile import load_locustfile
 
 
 class _HeaderCheckMixin:
@@ -656,6 +659,41 @@ class TestWebUI(LocustTestCase, _HeaderCheckMixin):
         response = requests.get("http://127.0.0.1:%i/stop" % self.web_port)
         self.assertEqual(response.json()["message"], "Test stopped")
 
+    def test_swarm_shape_class_is_updated_when_userclass_picker_is_active(self):
+        class User1(User):
+            pass
+
+        class TestShape(LoadTestShape):
+            def tick(self):
+                pass
+
+        test_shape_instance = TestShape()
+
+        self.environment.web_ui.userclass_picker_is_active = True
+        self.environment.available_user_classes = {"User1": User1}
+        self.environment.available_shape_classes = {"TestShape": test_shape_instance}
+        self.environment.shape_class = None
+
+        response = requests.post(
+            "http://127.0.0.1:%i/swarm" % self.web_port,
+            data={
+                "user_count": 5,
+                "spawn_rate": 5,
+                "host": "https://localhost",
+                "user_classes": "User1",
+                "shape_class": "TestShape",
+            },
+        )
+
+        self.assertEqual(200, response.status_code)
+        self.assertEqual(test_shape_instance, self.environment.shape_class)
+        self.assertIsNotNone(test_shape_instance.runner)
+
+        # stop
+        gevent.sleep(1)
+        response = requests.get("http://127.0.0.1:%i/stop" % self.web_port)
+        self.assertEqual(response.json()["message"], "Test stopped")
+
     def test_swarm_userclass_shapeclass_ignored_when_userclass_picker_is_inactive(self):
         class User1(User):
             wait_time = constant(1)
@@ -1027,17 +1065,76 @@ class TestWebUI(LocustTestCase, _HeaderCheckMixin):
 
         self.assertIn(log_line, response.json().get("logs"))
 
+    def test_template_args(self):
+        class MyUser(User):
+            @task
+            def do_something(self):
+                self.client.get("/")
+
+            host = "http://example.com"
+
+        class MyUser2(User):
+            host = "http://example.com"
+
+        self.environment.user_classes = [MyUser, MyUser2]
+        self.environment.available_user_classes = {"User1": MyUser, "User2": MyUser2}
+        self.environment.available_user_tasks = {"User1": MyUser.tasks, "User2": MyUser2.tasks}
+
+        users = {"User1": MyUser.json(), "User2": MyUser2.json()}
+        available_user_tasks = {"User1": ["do_something"], "User2": []}
+
+        self.web_ui.update_template_args()
+
+        self.assertEqual(self.web_ui.template_args.get("users"), users)
+        self.assertEqual(
+            self.web_ui.template_args.get("available_user_classes"), sorted(self.environment.available_user_classes)
+        )
+        self.assertEqual(self.web_ui.template_args.get("available_user_tasks"), available_user_tasks)
+
+    def test_update_user_endpoint(self):
+        class MyUser(User):
+            @task
+            def my_task(self):
+                pass
+
+            @task
+            def my_task_2(self):
+                pass
+
+            host = "http://example.com"
+
+        class MyUser2(User):
+            host = "http://example.com"
+
+        self.environment.user_classes = [MyUser, MyUser2]
+        self.environment.available_user_classes = {"User1": MyUser, "User2": MyUser2}
+        self.environment.available_user_tasks = {"User1": MyUser.tasks, "User2": MyUser2.tasks}
+
+        users = {"User1": MyUser.json(), "User2": MyUser2.json()}
+        available_user_tasks = {"User1": ["my_task", "my_task_2"], "User2": []}
+
+        # environment.update_user_class({"user_class_name": "User1", "host": "http://localhost", "tasks": ["my_task_2"]})
+        response = requests.post(
+            "http://127.0.0.1:%i/user" % self.web_port,
+            json={"user_class_name": "User1", "host": "http://localhost", "tasks": ["my_task_2"]},
+        )
+
+        self.assertEqual(
+            self.environment.available_user_classes["User1"].json(),
+            {"host": "http://localhost", "tasks": ["my_task_2"], "fixed_count": 0, "weight": 1},
+        )
+
 
 class TestWebUIAuth(LocustTestCase):
     def setUp(self):
         super().setUp()
 
         parser = get_parser(default_config_files=[])
-        options = parser.parse_args(["--web-auth", "john:doe"])
-        self.runner = Runner(self.environment)
-        self.stats = self.runner.stats
-        self.web_ui = self.environment.create_web_ui("127.0.0.1", 0, auth_credentials=options.web_auth)
-        self.web_ui.app.view_functions["request_stats"].clear_cache()
+        self.environment.parsed_options = parser.parse_args(["--modern-ui", "--web-login"])
+
+        self.web_ui = self.environment.create_web_ui("127.0.0.1", 0, modern_ui=True, web_login=True)
+
+        self.web_ui.app.secret_key = "secret!"
         gevent.sleep(0.01)
         self.web_port = self.web_ui.server.server_port
 
@@ -1046,18 +1143,36 @@ class TestWebUIAuth(LocustTestCase):
         self.web_ui.stop()
         self.runner.quit()
 
-    def test_index_with_basic_auth_enabled_correct_credentials(self):
-        self.assertEqual(
-            200, requests.get("http://127.0.0.1:%i/?ele=phino" % self.web_port, auth=("john", "doe")).status_code
-        )
+    def test_index_with_web_login_enabled_valid_user(self):
+        class User(UserMixin):
+            def __init__(self):
+                self.username = "test_user"
 
-    def test_index_with_basic_auth_enabled_incorrect_credentials(self):
-        self.assertEqual(
-            401, requests.get("http://127.0.0.1:%i/?ele=phino" % self.web_port, auth=("john", "invalid")).status_code
-        )
+            def get_id(self):
+                return self.username
+
+        def load_user(id):
+            return User()
+
+        self.web_ui.login_manager.request_loader(load_user)
+
+        response = requests.get("http://127.0.0.1:%i" % self.web_port)
+        d = pq(response.content.decode("utf-8"))
+
+        self.assertNotIn("authArgs", str(d))
+        self.assertIn("templateArgs", str(d))
+
+    def test_index_with_web_login_enabled_no_user(self):
+        def load_user():
+            return None
+
+        self.web_ui.login_manager.user_loader(load_user)
+
+        response = requests.get("http://127.0.0.1:%i" % self.web_port)
+        d = pq(response.content.decode("utf-8"))
 
-    def test_index_with_basic_auth_enabled_blank_credentials(self):
-        self.assertEqual(401, requests.get("http://127.0.0.1:%i/?ele=phino" % self.web_port).status_code)
+        # asserts auth page is returned
+        self.assertIn("authArgs", str(d))
 
 
 class TestWebUIWithTLS(LocustTestCase):

locust/test/test_zmqrpc.py

@@ -1,8 +1,10 @@
+from locust.exception import RPCError, RPCReceiveError, RPCSendError
+from locust.rpc import Message, zmqrpc
+from locust.test.testcases import LocustTestCase
+
 from time import sleep
+
 import zmq
-from locust.rpc import zmqrpc, Message
-from locust.test.testcases import LocustTestCase
-from locust.exception import RPCError, RPCSendError, RPCReceiveError
 
 
 class ZMQRPC_tests(LocustTestCase):

locust/test/testcases.py

@@ -1,3 +1,10 @@
+import locust
+from locust import log
+from locust.env import Environment
+from locust.event import Events
+from locust.test.mock_logging import MockedLoggingHandler
+from locust.test.util import clear_all_functools_lru_cache
+
 import base64
 import logging
 import random
@@ -10,13 +17,6 @@ import gevent
 import gevent.pywsgi
 from flask import Flask, Response, make_response, redirect, request, send_file, stream_with_context
 
-import locust
-from locust import log
-from locust.event import Events
-from locust.env import Environment
-from locust.test.mock_logging import MockedLoggingHandler
-from locust.test.util import clear_all_functools_lru_cache
-
 app = Flask(__name__)
 app.jinja_env.add_extension("jinja2.ext.do")
 

locust/test/util.py

@@ -1,17 +1,16 @@
+import datetime
 import functools
 import gc
 import os
 import socket
-import datetime
+from contextlib import contextmanager
+from tempfile import NamedTemporaryFile
+
 from cryptography import x509
-from cryptography.x509.oid import NameOID
-from cryptography.hazmat.primitives import hashes
 from cryptography.hazmat.backends import default_backend
-from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives import hashes, serialization
 from cryptography.hazmat.primitives.asymmetric import rsa
-
-from contextlib import contextmanager
-from tempfile import NamedTemporaryFile
+from cryptography.x509.oid import NameOID
 
 
 @contextmanager

locust/user/__init__.py

@@ -1,2 +1,2 @@
-from .task import task, tag, TaskSet
+from .task import TaskSet, tag, task
 from .users import HttpUser, User