locust 2.18.1.dev4__py3-none-any.whl → 2.18.1.dev16__py3-none-any.whl

locust/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '2.18.1.dev4'
-__version_tuple__ = version_tuple = (2, 18, 1, 'dev4')
+__version__ = version = '2.18.1.dev16'
+__version_tuple__ = version_tuple = (2, 18, 1, 'dev16')

locust/log.py

@@ -9,6 +9,15 @@ HOSTNAME = socket.gethostname()
 unhandled_greenlet_exception = False
 
 
+class LogReader(logging.Handler):
+    def __init__(self):
+        super().__init__()
+        self.logs = []
+
+    def emit(self, record):
+        self.logs.append(self.format(record))
+
+
 def setup_logging(loglevel, logfile=None):
     loglevel = loglevel.upper()
 
@@ -32,21 +41,22 @@ def setup_logging(loglevel, logfile=None):
                 "class": "logging.StreamHandler",
                 "formatter": "plain",
             },
+            "log_reader": {"class": "locust.log.LogReader", "formatter": "default"},
         },
         "loggers": {
             "locust": {
-                "handlers": ["console"],
+                "handlers": ["console", "log_reader"],
                 "level": loglevel,
                 "propagate": False,
             },
             "locust.stats_logger": {
-                "handlers": ["console_plain"],
+                "handlers": ["console_plain", "log_reader"],
                 "level": "INFO",
                 "propagate": False,
             },
         },
         "root": {
-            "handlers": ["console"],
+            "handlers": ["console", "log_reader"],
             "level": loglevel,
         },
     }
@@ -58,8 +68,8 @@ def setup_logging(loglevel, logfile=None):
             "filename": logfile,
             "formatter": "default",
         }
-        LOGGING_CONFIG["loggers"]["locust"]["handlers"] = ["file"]
-        LOGGING_CONFIG["root"]["handlers"] = ["file"]
+        LOGGING_CONFIG["loggers"]["locust"]["handlers"] = ["file", "log_reader"]
+        LOGGING_CONFIG["root"]["handlers"] = ["file", "log_reader"]
 
     logging.config.dictConfig(LOGGING_CONFIG)
 

locust/main.py

@@ -86,17 +86,20 @@ def main():
     user_classes: Dict[str, locust.User] = {}
     available_user_classes = {}
     available_shape_classes = {}
+    shape_class = None
     for _locustfile in locustfiles:
-        docstring, _user_classes, shape_class = load_locustfile(_locustfile)
+        docstring, _user_classes, shape_classes = load_locustfile(_locustfile)
 
         # Setting Available Shape Classes
-        if shape_class:
-            shape_class_name = type(shape_class).__name__
-            if shape_class_name in available_shape_classes.keys():
-                sys.stderr.write(f"Duplicate shape classes: {shape_class_name}\n")
-                sys.exit(1)
+        if shape_classes:
+            shape_class = shape_classes[0]
+            for shape_class in shape_classes:
+                shape_class_name = type(shape_class).__name__
+                if shape_class_name in available_shape_classes.keys():
+                    sys.stderr.write(f"Duplicate shape classes: {shape_class_name}\n")
+                    sys.exit(1)
 
-            available_shape_classes[shape_class_name] = shape_class
+                available_shape_classes[shape_class_name] = shape_class
 
         # Setting Available User Classes
         for key, value in _user_classes.items():

locust/static/css/application.css

@@ -470,5 +470,22 @@ body.ready .main {
   margin-right: 10px;
   color: #48a584;
 }
+.userClass {
+  border: none;
+  background: #fff;
+  height: 100px;
+  width: 340px;
+  font-size: 18px;
+  padding-left: 10px;
+}
+
+.shapeClass {
+  border: none;
+  background: #fff;
+  /* height: 100px; */
+  width: 340px;
+  font-size: 18px;
+  padding-left: 10px;
+}
 
 /*# sourceMappingURL=application.css.map */

locust/templates/index.html

@@ -65,7 +65,7 @@
                         {% for user in available_user_classes %}
                             <option value="{{user}}" selected>{{user}}</option>
                         {% endfor %}
-                        </select><br>range
+                        </select><br>
                         <label for="available_shape_classes">ShapeClass </label>
                         <select name="shape_class" id="shape-classes" class="shapeClass">
                         {% for shape_class in available_shape_classes %}

locust/test/test_load_locustfile.py

@@ -37,25 +37,25 @@ class TestLoadLocustfile(LocustTestCase):
 
     def test_load_locust_file_from_absolute_path(self):
         with mock_locustfile() as mocked:
-            docstring, user_classes, shape_class = main.load_locustfile(mocked.file_path)
+            docstring, user_classes, shape_classes = main.load_locustfile(mocked.file_path)
             self.assertIn("UserSubclass", user_classes)
             self.assertNotIn("NotUserSubclass", user_classes)
             self.assertNotIn("LoadTestShape", user_classes)
-            self.assertIsNone(shape_class)
+            self.assertEqual(shape_classes, [])
 
     def test_load_locust_file_from_relative_path(self):
         with mock_locustfile() as mocked:
-            docstring, user_classes, shape_class = main.load_locustfile(
+            docstring, user_classes, shape_classes = main.load_locustfile(
                 os.path.join(os.path.relpath(mocked.directory, os.getcwd()), mocked.filename)
             )
 
     def test_load_locust_file_with_a_dot_in_filename(self):
         with mock_locustfile(filename_prefix="mocked.locust.file") as mocked:
-            docstring, user_classes, shape_class = main.load_locustfile(mocked.file_path)
+            docstring, user_classes, shape_classes = main.load_locustfile(mocked.file_path)
 
     def test_return_docstring_and_user_classes(self):
         with mock_locustfile() as mocked:
-            docstring, user_classes, shape_class = main.load_locustfile(mocked.file_path)
+            docstring, user_classes, shape_classes = main.load_locustfile(mocked.file_path)
             self.assertEqual("This is a mock locust file for unit testing", docstring)
             self.assertIn("UserSubclass", user_classes)
             self.assertNotIn("NotUserSubclass", user_classes)
@@ -70,11 +70,31 @@ class TestLoadLocustfile(LocustTestCase):
         """
         )
         with mock_locustfile(content=content) as mocked:
-            docstring, user_classes, shape_class = main.load_locustfile(mocked.file_path)
+            docstring, user_classes, shape_classes = main.load_locustfile(mocked.file_path)
             self.assertEqual("This is a mock locust file for unit testing", docstring)
             self.assertIn("UserSubclass", user_classes)
             self.assertNotIn("NotUserSubclass", user_classes)
-            self.assertEqual(shape_class.__class__.__name__, "LoadTestShape")
+            self.assertEqual(shape_classes[0].__class__.__name__, "LoadTestShape")
+
+    def test_with_multiple_shape_classes(self):
+        content = MOCK_LOCUSTFILE_CONTENT + textwrap.dedent(
+            """\
+        class LoadTestShape1(LoadTestShape):
+            def tick(self):
+                pass
+
+        class LoadTestShape2(LoadTestShape):
+            def tick(self):
+                pass
+        """
+        )
+        with mock_locustfile(content=content) as mocked:
+            docstring, user_classes, shape_classes = main.load_locustfile(mocked.file_path)
+            self.assertEqual("This is a mock locust file for unit testing", docstring)
+            self.assertIn("UserSubclass", user_classes)
+            self.assertNotIn("NotUserSubclass", user_classes)
+            self.assertEqual(shape_classes[0].__class__.__name__, "LoadTestShape1")
+            self.assertEqual(shape_classes[1].__class__.__name__, "LoadTestShape2")
 
     def test_with_abstract_shape_class(self):
         content = MOCK_LOCUSTFILE_CONTENT + textwrap.dedent(
@@ -92,10 +112,10 @@ class TestLoadLocustfile(LocustTestCase):
         )
 
         with mock_locustfile(content=content) as mocked:
-            _, user_classes, shape_class = main.load_locustfile(mocked.file_path)
+            _, user_classes, shape_classes = main.load_locustfile(mocked.file_path)
             self.assertNotIn("UserBaseLoadTestShape", user_classes)
             self.assertNotIn("UserLoadTestShape", user_classes)
-            self.assertEqual(shape_class.__class__.__name__, "UserLoadTestShape")
+            self.assertEqual(shape_classes[0].__class__.__name__, "UserLoadTestShape")
 
     def test_with_not_imported_shape_class(self):
         content = MOCK_LOCUSTFILE_CONTENT + textwrap.dedent(
@@ -107,9 +127,9 @@ class TestLoadLocustfile(LocustTestCase):
         )
 
         with mock_locustfile(content=content) as mocked:
-            _, user_classes, shape_class = main.load_locustfile(mocked.file_path)
+            _, user_classes, shape_classes = main.load_locustfile(mocked.file_path)
             self.assertNotIn("UserLoadTestShape", user_classes)
-            self.assertEqual(shape_class.__class__.__name__, "UserLoadTestShape")
+            self.assertEqual(shape_classes[0].__class__.__name__, "UserLoadTestShape")
 
     def test_create_environment(self):
         options = parse_options(

locust/test/test_web.py

@@ -5,13 +5,13 @@ import os
 import re
 import textwrap
 import traceback
+import logging
 from io import StringIO
 from tempfile import NamedTemporaryFile, TemporaryDirectory
 
 import gevent
 import requests
 from pyquery import PyQuery as pq
-
 import locust
 from locust import constant, LoadTestShape
 from locust.argument_parser import get_parser, parse_options
@@ -21,6 +21,7 @@ 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 .mock_locustfile import mock_locustfile
 from .testcases import LocustTestCase
@@ -563,7 +564,15 @@ class TestWebUI(LocustTestCase, _HeaderCheckMixin):
             def t(self):
                 pass
 
-        class TestShape(LoadTestShape):
+        class TestShape1(LoadTestShape):
+            def tick(self):
+                run_time = self.get_run_time()
+                if run_time < 10:
+                    return 4, 4
+                else:
+                    return None
+
+        class TestShape2(LoadTestShape):
             def tick(self):
                 run_time = self.get_run_time()
                 if run_time < 10:
@@ -573,8 +582,8 @@ class TestWebUI(LocustTestCase, _HeaderCheckMixin):
 
         self.environment.web_ui.userclass_picker_is_active = True
         self.environment.available_user_classes = {"User1": User1, "User2": User2}
-        self.environment.available_shape_classes = {"TestShape": TestShape()}
-        self.environment.shape_class = None
+        self.environment.available_shape_classes = {"TestShape1": TestShape1(), "TestShape2": TestShape2()}
+        self.environment.shape_class = TestShape1()
 
         response = requests.post(
             "http://127.0.0.1:%i/swarm" % self.web_port,
@@ -583,14 +592,14 @@ class TestWebUI(LocustTestCase, _HeaderCheckMixin):
                 "spawn_rate": 5,
                 "host": "https://localhost",
                 "user_classes": "User1",
-                "shape_class": "TestShape",
+                "shape_class": "TestShape2",
             },
         )
 
         self.assertEqual(200, response.status_code)
         self.assertEqual("https://localhost", response.json()["host"])
         self.assertEqual(self.environment.host, "https://localhost")
-        assert isinstance(self.environment.shape_class, TestShape)
+        assert isinstance(self.environment.shape_class, TestShape2)
 
         # stop
         gevent.sleep(1)
@@ -1005,6 +1014,19 @@ class TestWebUI(LocustTestCase, _HeaderCheckMixin):
         self.assertIn("Script: <span>locust.py</span>", str(d))
         self.assertIn("Target Host: <span>http://localhost</span>", str(d))
 
+    def test_logs(self):
+        log_handler = LogReader()
+        log_handler.name = "log_reader"
+        log_handler.setLevel(logging.INFO)
+        logger = logging.getLogger("root")
+        logger.addHandler(log_handler)
+        log_line = "some log info"
+        logger.info(log_line)
+
+        response = requests.get("http://127.0.0.1:%i/logs" % self.web_port)
+
+        self.assertIn(log_line, response.json().get("logs"))
+
 
 class TestWebUIAuth(LocustTestCase):
     def setUp(self):

locust/util/load_locustfile.py

@@ -2,7 +2,7 @@ import importlib
 import inspect
 import os
 import sys
-from typing import Dict, Optional, Tuple
+from typing import Dict, List, Optional, Tuple
 from ..shape import LoadTestShape
 from ..user import User
 
@@ -21,7 +21,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], Optional[LoadTestShape]]:
+def load_locustfile(path) -> Tuple[Optional[str], Dict[str, User], List[LoadTestShape]]:
     """
     Import given locustfile path and return (docstring, callables).
 
@@ -65,10 +65,6 @@ def load_locustfile(path) -> Tuple[Optional[str], Dict[str, User], Optional[Load
     user_classes = {name: value for name, value in vars(imported).items() if is_user_class(value)}
 
     # Find shape class, if any, return it
-    shape_classes = [value for name, value in vars(imported).items() if is_shape_class(value)]
-    if shape_classes:
-        shape_class = shape_classes[0]()
-    else:
-        shape_class = None
+    shape_classes = [value() for name, value in vars(imported).items() if is_shape_class(value)]
 
-    return imported.__doc__, user_classes, shape_class
+    return imported.__doc__, user_classes, shape_classes

locust/web.py

@@ -136,6 +136,12 @@ class WebUI:
         if not delayed_start:
             self.start()
 
+        @app.errorhandler(Exception)
+        def handle_exception(error):
+            error_message = str(error)
+            logger.log(logging.CRITICAL, error_message)
+            return make_response(error_message, 500)
+
         @app.route("/assets/<path:path>")
         def send_assets(path):
             webui_build_path = self.webui_build_path
@@ -229,7 +235,11 @@ class WebUI:
             if environment.shape_class and environment.runner is not None:
                 environment.runner.start_shape()
                 return jsonify(
-                    {"success": True, "message": "Swarming started using shape class", "host": environment.host}
+                    {
+                        "success": True,
+                        "message": f"Swarming started using shape class '{type(environment.shape_class).__name__}'",
+                        "host": environment.host,
+                    }
                 )
 
             if self._swarm_greenlet is not None:
@@ -474,6 +484,20 @@ class WebUI:
             }
             return task_data
 
+        @app.route("/logs")
+        @self.auth_required_if_enabled
+        def logs():
+            log_reader_handler = [
+                handler for handler in logging.getLogger("root").handlers if handler.name == "log_reader"
+            ]
+
+            if log_reader_handler:
+                logs = log_reader_handler[0].logs
+            else:
+                logs = []
+
+            return jsonify({"logs": logs})
+
     def start(self):
         self.greenlet = gevent.spawn(self.start_server)
         self.greenlet.link_exception(greenlet_exception_handler)