locust 2.29.2.dev8__tar.gz → 2.29.2.dev15__tar.gz

{locust-2.29.2.dev8 → locust-2.29.2.dev15}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: locust
-Version: 2.29.2.dev8
+Version: 2.29.2.dev15
 Summary: Developer-friendly load testing framework
 License: MIT
 Project-URL: Homepage, https://github.com/locustio/locust

{locust-2.29.2.dev8 → locust-2.29.2.dev15}/examples/test_data_management.py

@@ -6,6 +6,10 @@
 # 3. Test start
 # 4. User start
 # 5. Inside a task
+# M1. CPU & memory usage
+# M2. master sent heartbeat to worker 1-N
+# M3. worker 1-N received heartbeat from master
+# (M* are repeated as long as locust is running)
 # ...
 # 6. Test run stopping
 # 7. User stop
@@ -15,17 +19,21 @@
 # 10. Locust quit
 #
 # try it out by running:
-#  locust -f test_data_management.py --headless -u 2 -t 5
+#  locust -f test_data_management.py --headless -u 2 -t 5 --processes 2
+from __future__ import annotations
+
 from locust import HttpUser, events, task
+from locust.env import Environment
 from locust.runners import MasterRunner
 from locust.user.wait_time import constant
 
 import datetime
+from typing import Any
 
 import requests
 
 
-def timestring():
+def timestring() -> str:
     now = datetime.datetime.now()
     return datetime.datetime.strftime(now, "%m:%S.%f")[:-5]
 
@@ -42,17 +50,17 @@ test_run_specific_data = None
 
 
 @events.init.add_listener
-def init(environment, **_kwargs):
+def init(environment: Environment, **_kwargs: Any) -> None:
     print("2. Initializing locust, happens after parsing the locustfile but before test start")
 
 
 @events.quitting.add_listener
-def quitting(environment, **_kwargs):
+def quitting(environment: Environment, **_kwargs: Any) -> None:
     print("9. locust is about to shut down")
 
 
 @events.test_start.add_listener
-def test_start(environment, **_kwargs):
+def test_start(environment: Environment, **_kwargs) -> None:
     # happens only once in headless runs, but can happen multiple times in web ui-runs
     global test_run_specific_data
     print("3. Starting test run")
@@ -64,18 +72,35 @@ def test_start(environment, **_kwargs):
         ).json()["data"]
 
 
+@events.heartbeat_sent.add_listener
+def heartbeat_sent(client_id: str, timestamp: float) -> None:
+    print(f"M2. master sent heartbeat to worker {client_id} at {datetime.datetime.fromtimestamp(timestamp)}")
+
+
+@events.heartbeat_received.add_listener
+def heartbeat_received(client_id: str, timestamp: float) -> None:
+    print(f"M3. worker {client_id} received heartbeat from master at {datetime.datetime.fromtimestamp(timestamp)}")
+
+
+@events.usage_monitor.add_listener
+def usage_monitor(environment: Environment, cpu_usage: float, memory_usage: int) -> None:
+    # convert from bytes to Mebibytes
+    memory_usage = memory_usage / 1024 / 1024
+    print(f"M1. {environment.runner.__class__.__name__}: cpu={cpu_usage}%, memory={memory_usage}M")
+
+
 @events.quit.add_listener
-def quit(exit_code, **kwargs):
+def quit(exit_code: int, **kwargs: Any) -> None:
     print(f"10. Locust has shut down with code {exit_code}")
 
 
 @events.test_stopping.add_listener
-def test_stopping(environment, **_kwargs):
+def test_stopping(environment: Environment, **_kwargs: Any) -> None:
     print("6. stopping test run")
 
 
 @events.test_stop.add_listener
-def test_stop(environment, **_kwargs):
+def test_stop(environment: Environment, **_kwargs: Any) -> None:
     print("8. test run stopped")
 
 
@@ -84,7 +109,7 @@ class MyUser(HttpUser):
     wait_time = constant(180)  # be nice to postman-echo
     first_start = True
 
-    def on_start(self):
+    def on_start(self) -> None:
         if MyUser.first_start:
             MyUser.first_start = False
             # This is useful for similar things as to test_start, but happens in the context of a User
@@ -101,7 +126,7 @@ class MyUser(HttpUser):
         ).json()["data"]
 
     @task
-    def t(self):
+    def t(self) -> None:
         self.client.get(f"/get?{global_test_data}")
         self.client.get(f"/get?{test_run_specific_data}")
         self.client.get(f"/get?{self.user_specific_testdata}")
@@ -114,6 +139,6 @@ class MyUser(HttpUser):
         ).json()["data"]
         self.client.get(f"/get?{task_run_specific_testdata}")
 
-    def on_stop(self):
+    def on_stop(self) -> None:
         # this is a good place to clean up/release any user-specific test data
         print("7. a user was stopped")

{locust-2.29.2.dev8 → locust-2.29.2.dev15}/locust/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '2.29.2.dev8'
-__version_tuple__ = version_tuple = (2, 29, 2, 'dev8')
+__version__ = version = '2.29.2.dev15'
+__version_tuple__ = version_tuple = (2, 29, 2, 'dev15')

{locust-2.29.2.dev8 → locust-2.29.2.dev15}/locust/contrib/fasthttp.py

@@ -12,12 +12,11 @@ import socket
 import time
 import traceback
 from base64 import b64encode
-from collections.abc import Generator
 from contextlib import contextmanager
 from http.cookiejar import CookieJar
 from json.decoder import JSONDecodeError
 from ssl import SSLError
-from typing import Any, Callable, cast
+from typing import TYPE_CHECKING, cast
 from urllib.parse import urlparse, urlunparse
 
 import gevent
@@ -32,6 +31,36 @@ from geventhttpclient.useragent import CompatRequest, CompatResponse, Connection
 # borrow requests's content-type header parsing
 from requests.utils import get_encoding_from_headers
 
+if TYPE_CHECKING:
+    import sys
+    from collections.abc import Callable, Generator
+    from typing import TypedDict
+
+    if sys.version_info >= (3, 11):
+        from typing import Unpack
+    else:
+        from typing_extensions import Unpack
+
+    class PostKwargs(TypedDict, total=False):
+        name: str | None
+        catch_response: bool
+        stream: bool
+        headers: dict | None
+        auth: tuple[str | bytes, str | bytes] | None
+        allow_redirects: bool
+        context: dict
+
+    class PutKwargs(PostKwargs, total=False):
+        json: dict | None
+
+    class PatchKwargs(PostKwargs, total=False):
+        json: dict | None
+
+    class RESTKwargs(PostKwargs, total=False):
+        data: str | dict | None
+        json: dict | None
+
+
 # Monkey patch geventhttpclient.useragent.CompatRequest so that Cookiejar works with Python >= 3.3
 # More info: https://github.com/requests/requests/pull/871
 CompatRequest.unverifiable = False
@@ -85,7 +114,7 @@ class FastHttpSession:
         client_pool: HTTPClientPool | None = None,
         ssl_context_factory: Callable | None = None,
         **kwargs,
-    ):
+    ) -> None:
         self.environment = environment
         self.base_url = base_url
         self.cookiejar = CookieJar()
@@ -117,14 +146,14 @@ class FastHttpSession:
             # store authentication header (we construct this by using _basic_auth_str() function from requests.auth)
             self.auth_header = _construct_basic_auth_str(parsed_url.username, parsed_url.password)
 
-    def _build_url(self, path):
+    def _build_url(self, path: str) -> str:
         """prepend url with hostname unless it's already an absolute URL"""
         if absolute_http_url_regexp.match(path):
             return path
         else:
             return f"{self.base_url}{path}"
 
-    def _send_request_safe_mode(self, method, url, **kwargs):
+    def _send_request_safe_mode(self, method: str, url: str, **kwargs):
         """
         Send an HTTP request, and catch any exception that might occur due to either
         connection problems, or invalid HTTP status codes
@@ -155,7 +184,7 @@ class FastHttpSession:
         catch_response: bool = False,
         stream: bool = False,
         headers: dict | None = None,
-        auth=None,
+        auth: tuple[str | bytes, str | bytes] | None = None,
         json: dict | None = None,
         allow_redirects: bool = True,
         context: dict = {},
@@ -264,31 +293,37 @@ class FastHttpSession:
             self.environment.events.request.fire(**request_meta)
             return response
 
-    def delete(self, url, **kwargs):
+    def delete(self, url: str, **kwargs: Unpack[RESTKwargs]) -> ResponseContextManager | FastResponse:
         """Sends a DELETE request"""
         return self.request("DELETE", url, **kwargs)
 
-    def get(self, url, **kwargs):
+    def get(self, url: str, **kwargs: Unpack[RESTKwargs]) -> ResponseContextManager | FastResponse:
         """Sends a GET request"""
         return self.request("GET", url, **kwargs)
 
-    def head(self, url, **kwargs):
+    def head(self, url: str, **kwargs: Unpack[RESTKwargs]) -> ResponseContextManager | FastResponse:
         """Sends a HEAD request"""
         return self.request("HEAD", url, **kwargs)
 
-    def options(self, url, **kwargs):
+    def options(self, url: str, **kwargs: Unpack[RESTKwargs]) -> ResponseContextManager | FastResponse:
         """Sends a OPTIONS request"""
         return self.request("OPTIONS", url, **kwargs)
 
-    def patch(self, url, data=None, **kwargs):
+    def patch(
+        self, url: str, data: str | dict | None = None, **kwargs: Unpack[PatchKwargs]
+    ) -> ResponseContextManager | FastResponse:
         """Sends a PATCH request"""
         return self.request("PATCH", url, data=data, **kwargs)
 
-    def post(self, url, data=None, json=None, **kwargs):
+    def post(
+        self, url: str, data: str | dict | None = None, json: dict | None = None, **kwargs: Unpack[PostKwargs]
+    ) -> ResponseContextManager | FastResponse:
         """Sends a POST request"""
         return self.request("POST", url, data=data, json=json, **kwargs)
 
-    def put(self, url, data=None, **kwargs):
+    def put(
+        self, url: str, data: str | dict | None = None, **kwargs: Unpack[PutKwargs]
+    ) -> ResponseContextManager | FastResponse:
         """Sends a PUT request"""
         return self.request("PUT", url, data=data, **kwargs)
 

{locust-2.29.2.dev8 → locust-2.29.2.dev15}/locust/event.py

@@ -235,6 +235,38 @@ class Events:
     Fired when the CPU usage exceeds runners.CPU_WARNING_THRESHOLD (90% by default)
     """
 
+    heartbeat_sent: EventHook
+    """
+    Fired when a heartbeat is sent by master to a worker.
+
+    Event arguments:
+
+    :param client_id: worker client id
+    :param timestamp: time in seconds since the epoch (float) when the event occured
+    """
+
+    heartbeat_received: EventHook
+    """
+    Fired when a heartbeat is received by a worker from master.
+
+    Event arguments:
+
+    :param client_id: worker client id
+    :param timestamp: time in seconds since the epoch (float) when the event occured
+    """
+
+    usage_monitor: EventHook
+    """
+    Fired every runners.CPU_MONITOR_INTERVAL (5.0 seconds by default) with information about
+    current CPU and memory usage.
+
+    Event arguments:
+
+    :param environment: locust environment
+    :param cpu_usage: current CPU usage in percent
+    :param memory_usage: current memory usage (RSS) in bytes
+    """
+
     def __init__(self):
         # For backward compatibility use also values of class attributes
         for name, value in vars(type(self)).items():

{locust-2.29.2.dev8 → locust-2.29.2.dev15}/locust/runners.py

@@ -15,19 +15,9 @@ import traceback
 from abc import abstractmethod
 from collections import defaultdict
 from collections.abc import Iterator, MutableMapping, ValuesView
-from operator import (
-    itemgetter,
-    methodcaller,
-)
+from operator import itemgetter, methodcaller
 from types import TracebackType
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    NoReturn,
-    TypedDict,
-    cast,
-)
+from typing import TYPE_CHECKING, Any, Callable, NoReturn, TypedDict, cast
 from uuid import uuid4
 
 import gevent
@@ -40,15 +30,8 @@ from . import argument_parser
 from .dispatch import UsersDispatcher
 from .exception import RPCError, RPCReceiveError, RPCSendError
 from .log import get_logs, greenlet_exception_logger
-from .rpc import (
-    Message,
-    rpc,
-)
-from .stats import (
-    RequestStats,
-    StatsError,
-    setup_distributed_stats_event_listeners,
-)
+from .rpc import Message, rpc
+from .stats import RequestStats, StatsError, setup_distributed_stats_event_listeners
 
 if TYPE_CHECKING:
     from . import User
@@ -106,7 +89,7 @@ class Runner:
         self.spawning_greenlet: gevent.Greenlet | None = None
         self.shape_greenlet: gevent.Greenlet | None = None
         self.shape_last_tick: tuple[int, float] | tuple[int, float, list[type[User]] | None] | None = None
-        self.current_cpu_usage: int = 0
+        self.current_cpu_usage: float = 0.0
         self.cpu_warning_emitted: bool = False
         self.worker_cpu_warning_emitted: bool = False
         self.current_memory_usage: int = 0
@@ -308,6 +291,10 @@ class Runner:
                         f"CPU usage above {CPU_WARNING_THRESHOLD}%! This may constrain your throughput and may even give inconsistent response time measurements! See https://docs.locust.io/en/stable/running-distributed.html for how to distribute the load over multiple CPU cores or machines"
                     )
                     self.cpu_warning_emitted = True
+
+            self.environment.events.usage_monitor.fire(
+                environment=self.environment, cpu_usage=self.current_cpu_usage, memory_usage=self.current_memory_usage
+            )
             gevent.sleep(CPU_MONITOR_INTERVAL)
 
     @abstractmethod
@@ -1102,6 +1089,7 @@ class MasterRunner(DistributedRunner):
                         )
                     if "current_memory_usage" in msg.data:
                         c.memory_usage = msg.data["current_memory_usage"]
+                    self.environment.events.heartbeat_sent.fire(client_id=msg.node_id, timestamp=time.time())
                     self.server.send_to_client(Message("heartbeat", None, msg.node_id))
                 else:
                     logging.debug(f"Got heartbeat message from unknown worker {msg.node_id}")
@@ -1399,6 +1387,9 @@ class WorkerRunner(DistributedRunner):
                 self.reset_connection()
             elif msg.type == "heartbeat":
                 self.last_heartbeat_timestamp = time.time()
+                self.environment.events.heartbeat_received.fire(
+                    client_id=msg.node_id, timestamp=self.last_heartbeat_timestamp
+                )
             elif msg.type == "update_user_class":
                 self.environment.update_user_class(msg.data)
             elif msg.type == "spawning_complete":

{locust-2.29.2.dev8 → locust-2.29.2.dev15}/locust/test/test_runners.py

@@ -1,12 +1,7 @@
 from __future__ import annotations
 
 import locust
-from locust import (
-    LoadTestShape,
-    __version__,
-    constant,
-    runners,
-)
+from locust import LoadTestShape, __version__, constant, runners
 from locust.argument_parser import parse_options
 from locust.dispatch import UsersDispatcher
 from locust.env import Environment
@@ -26,11 +21,7 @@ from locust.runners import (
     WorkerRunner,
 )
 from locust.stats import RequestStats
-from locust.user import (
-    TaskSet,
-    User,
-    task,
-)
+from locust.user import TaskSet, User, task
 
 import json
 import logging
@@ -2136,6 +2127,149 @@ class TestMasterWorkerRunners(LocustTestCase):
 
             self.assertEqual(test_start_exec_count, 1)
 
+    def test_heartbeat_event(self) -> None:
+        """
+        Tests that heartbeat event is fired during a test
+        """
+
+        class TestUser(User):
+            wait_time = constant(0.1)
+
+            @task
+            def noop(self) -> None:
+                pass
+
+        with mock.patch("locust.runners.HEARTBEAT_INTERVAL", new=1):
+            # start a Master runner
+            master_env = Environment(user_classes=[TestUser])
+            worker_connect_events = []
+            timestamp_start: list[float] = [time.time() + 3600.0]
+
+            def on_connect(client_id: str) -> None:
+                worker_connect_events.append(client_id)
+                timestamp_start[0] = time.time()
+
+            master_env.events.worker_connect.add_listener(on_connect)
+            master = master_env.create_master_runner("*", 0)
+            sleep(0)
+            worker_env = Environment(user_classes=[TestUser])
+            worker: WorkerRunner = worker_env.create_worker_runner("127.0.0.1", master.server.port)
+
+            with (
+                mock.patch.object(
+                    worker.environment.events.heartbeat_received,
+                    "fire",
+                    wraps=worker.environment.events.heartbeat_received.fire,
+                ) as worker_heartbeat_received_mock,
+                mock.patch.object(
+                    master.environment.events.heartbeat_sent,
+                    "fire",
+                    wraps=master.environment.events.heartbeat_sent.fire,
+                ) as master_heartbeat_sent_mock,
+            ):
+                # give workers time to connect
+                sleep(0.1)
+                # issue start command that should trigger TestUsers to be spawned in the Workers
+                master.start(2, spawn_rate=2)
+                sleep(0.1)
+                # check that worker nodes have started locusts
+                self.assertEqual(2, worker.user_count)
+
+                # give time for nodes to send and receive 5 heartbeats, HEARTBEAT_INTERVAL mocked to 1 second, so
+                # sleep 5 seconds - 1 second that represents the overhead from connecting
+                sleep(5 - 1)
+                master.quit()
+
+                # make sure users are killed
+                self.assertEqual(0, worker.user_count)
+                # make sure events happened correctly
+                self.assertIn(worker.client_id, worker_connect_events)
+
+                timestamp_stop = time.time()
+
+                self.assertEqual(worker_heartbeat_received_mock.call_count, 5)
+                self.assertEqual(master_heartbeat_sent_mock.call_count, 5)
+
+                for call_args, call_kwargs in [
+                    *worker_heartbeat_received_mock.call_args_list,
+                    *master_heartbeat_sent_mock.call_args_list,
+                ]:
+                    self.assertEqual(call_args, ())  # args
+                    self.assertEqual(call_kwargs, {"client_id": worker.client_id, "timestamp": mock.ANY})  # kwargs
+                    self.assertGreaterEqual(call_kwargs["timestamp"], timestamp_start[0])
+                    self.assertLessEqual(call_kwargs["timestamp"], timestamp_stop)
+
+    def test_usage_monitor_event(self) -> None:
+        """
+        Tests that usage_monitor event is fired during a test
+        """
+
+        class TestUser(User):
+            wait_time = constant(0.1)
+
+            @task
+            def noop(self) -> None:
+                pass
+
+        with mock.patch("locust.runners.CPU_MONITOR_INTERVAL", new=1):
+            # start a Master runner
+            master_env = Environment(user_classes=[TestUser])
+            worker_connect_events = []
+
+            def on_connect(client_id: str) -> None:
+                worker_connect_events.append(client_id)
+
+            master_env.events.worker_connect.add_listener(on_connect)
+            master = master_env.create_master_runner("*", 0)
+            sleep(0)
+            worker_env = Environment(user_classes=[TestUser])
+            worker: WorkerRunner = worker_env.create_worker_runner("127.0.0.1", master.server.port)
+
+            with (
+                mock.patch.object(
+                    worker.environment.events.usage_monitor, "fire", wraps=worker.environment.events.usage_monitor.fire
+                ) as worker_usage_monitor_mock,
+                mock.patch.object(
+                    master.environment.events.usage_monitor, "fire", wraps=master.environment.events.usage_monitor.fire
+                ) as master_usage_monitor_mock,
+            ):
+                # give workers time to connect
+                sleep(0.1)
+                # issue start command that should trigger TestUsers to be spawned in the Workers
+                master.start(2, spawn_rate=2)
+                sleep(0.1)
+                # check that worker nodes have started locusts
+                self.assertEqual(2, worker.user_count)
+
+                # give time for nodes to send 5 usage_monitor events, CPU_MONITOR_INTERVAL mocked to 1 second, so
+                # sleep 5 seconds
+                sleep(5)
+                master.quit()
+
+                # make sure users are killed
+                self.assertEqual(0, worker.user_count)
+                # make sure events happened correctly
+                self.assertIn(worker.client_id, worker_connect_events)
+
+                self.assertEqual(worker_usage_monitor_mock.call_count, 5)
+                self.assertEqual(master_usage_monitor_mock.call_count, 5)
+
+                for call_args, call_kwargs in master_usage_monitor_mock:
+                    self.assertEqual(call_args, ())  # args
+                    self.assertEqual(
+                        call_kwargs, {"environment": master_env, "cpu_usage": mock.ANY, "memory_usage": mock.ANY}
+                    )  # kwargs
+                    self.assertTrue(isinstance(call_kwargs["cpu_usage"], float))
+                    self.assertTrue(isinstance(call_kwargs["memory_usage"], int))
+
+                for call_args, call_kwargs in worker_usage_monitor_mock:
+                    self.assertEqual(call_args, ())  # args
+                    self.assertEqual(
+                        call_kwargs, {"environment": worker_env, "cpu_usage": mock.ANY, "memory_usage": mock.ANY}
+                    )  # kwargs
+                    self.assertTrue(isinstance(call_kwargs["cpu_usage"], float))
+                    self.assertTrue(isinstance(call_kwargs["memory_usage"], int))
+
 
 class TestMasterRunner(LocustRunnerTestCase):
     def setUp(self):

{locust-2.29.2.dev8 → locust-2.29.2.dev15}/locust.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: locust
-Version: 2.29.2.dev8
+Version: 2.29.2.dev15
 Summary: Developer-friendly load testing framework
 License: MIT
 Project-URL: Homepage, https://github.com/locustio/locust