Merge pull request #2766 from andrewbaldwin44/feature/url-file-distri…

…bution

URL Directory, and Multi-File Support for File Distribution

locust/argument_parser.py

@@ -15,7 +15,6 @@
 import textwrap
 from collections import OrderedDict
 from typing import Any, NamedTuple
-from urllib.parse import urlparse
 from uuid import uuid4
 
 if sys.version_info >= (3, 11):
@@ -27,6 +26,9 @@
 import gevent
 import requests
 
+from .util.directory import get_abspaths_in
+from .util.url import is_url
+
 version = locust.__version__
 
 
@@ -125,14 +127,7 @@ def _parse_locustfile_path(path: str) -> list[str]:
         parsed_paths.append(download_locustfile_from_url(path))
     elif os.path.isdir(path):
         # Find all .py files in directory tree
-        for root, _dirs, fs in os.walk(path):
-            parsed_paths.extend(
-                [
-                    os.path.abspath(os.path.join(root, f))
-                    for f in fs
-                    if os.path.isfile(os.path.join(root, f)) and f.endswith(".py") and not f.startswith("_")
-                ]
-            )
+        parsed_paths.extend(get_abspaths_in(path, extension=".py"))
         if not parsed_paths:
             sys.stderr.write(f"Could not find any locustfiles in directory '{path}'")
             sys.exit(1)
@@ -148,20 +143,6 @@ def _parse_locustfile_path(path: str) -> list[str]:
     return parsed_paths
 
 
-def is_url(url: str) -> bool:
-    """
-    Check if path is an url
-    """
-    try:
-        result = urlparse(url)
-        if result.scheme == "https" or result.scheme == "http":
-            return True
-        else:
-            return False
-    except ValueError:
-        return False
-
-
 def download_locustfile_from_url(url: str) -> str:
     """
     Attempt to download and save locustfile from url.
@@ -238,7 +219,7 @@ def download_locustfile_from_master(master_host: str, master_port: int) -> str:
 
     def ask_for_locustfile():
         while not got_reply:
-            tempclient.send(Message("locustfile", None, client_id))
+            tempclient.send(Message("locustfile", {"version": version}, client_id))
             gevent.sleep(1)
 
     def log_warning():
@@ -271,14 +252,26 @@ def wait_for_reply():
         sys.stderr.write(f"Got error from master: {msg.data['error']}\n")
         sys.exit(1)
 
-    filename = msg.data["filename"]
-    with open(os.path.join(tempfile.gettempdir(), filename), "w", encoding="utf-8") as locustfile:
-        locustfile.write(msg.data["contents"])
+    tempclient.close()
+    return msg.data.get("locustfiles", [])
+
 
-    atexit.register(exit_handler, locustfile.name)
+def parse_locustfiles_from_master(locustfile_sources) -> list[str]:
+    locustfiles = []
 
-    tempclient.close()
-    return locustfile.name
+    for source in locustfile_sources:
+        if "contents" in source:
+            filename = source["filename"]
+            file_contents = source["contents"]
+
+            with open(os.path.join(tempfile.gettempdir(), filename), "w", encoding="utf-8") as locustfile:
+                locustfile.write(file_contents)
+
+            locustfiles.append(locustfile.name)
+        else:
+            locustfiles.append(source)
+
+    return locustfiles
 
 
 def parse_locustfile_option(args=None) -> list[str]:
@@ -339,10 +332,11 @@ def parse_locustfile_option(args=None) -> list[str]:
             )
             sys.exit(1)
         # having this in argument_parser module is a bit weird, but it needs to be done early
-        filename = download_locustfile_from_master(options.master_host, options.master_port)
-        return [filename]
+        locustfile_sources = download_locustfile_from_master(options.master_host, options.master_port)
+        locustfile_list = parse_locustfiles_from_master(locustfile_sources)
+    else:
+        locustfile_list = [f.strip() for f in options.locustfile.split(",")]
 
-    locustfile_list = [f.strip() for f in options.locustfile.split(",")]
     parsed_paths = parse_locustfile_paths(locustfile_list)
 
     if not parsed_paths:

locust/env.py

@@ -33,6 +33,7 @@ def __init__(
         stop_timeout: float | None = None,
         catch_exceptions=True,
         parsed_options: Namespace | None = None,
+        parsed_locustfiles: list[str] | 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,
@@ -91,6 +92,8 @@ def __init__(
         """
         self.parsed_options = parsed_options
         """Reference to the parsed command line options (used to pre-populate fields in Web UI). When using Locust as a library, this should either be `None` or an object created by `argument_parser.parse_args()`"""
+        self.parsed_locustfiles = parsed_locustfiles
+        """A list of all locustfiles for the test"""
         self.available_user_classes = available_user_classes
         """List of the available User Classes to pick from in the UserClass Picker"""
         self.available_shape_classes = available_shape_classes

locust/main.py

@@ -59,6 +59,7 @@ def create_environment(
     events=None,
     shape_class=None,
     locustfile=None,
+    parsed_locustfiles=None,
     available_user_classes=None,
     available_shape_classes=None,
     available_user_tasks=None,
@@ -74,6 +75,7 @@ def create_environment(
         host=options.host,
         reset_stats=options.reset_stats,
         parsed_options=options,
+        parsed_locustfiles=parsed_locustfiles,
         available_user_classes=available_user_classes,
         available_shape_classes=available_shape_classes,
         available_user_tasks=available_user_tasks,
@@ -349,6 +351,7 @@ def kill_workers(children):
         events=locust.events,
         shape_class=shape_class,
         locustfile=locustfile_path,
+        parsed_locustfiles=locustfiles,
         available_user_classes=available_user_classes,
         available_shape_classes=available_shape_classes,
         available_user_tasks=available_user_tasks,

locust/runners.py

@@ -32,6 +32,8 @@
 from .log import get_logs, greenlet_exception_logger
 from .rpc import Message, rpc
 from .stats import RequestStats, StatsError, setup_distributed_stats_event_listeners
+from .util.directory import get_abspaths_in
+from .util.url import is_url
 
 if TYPE_CHECKING:
     from . import User
@@ -1024,33 +1026,45 @@ def client_listener(self) -> NoReturn:
                 # if abs(time() - msg.data["time"]) > 5.0:
                 #    warnings.warn("The worker node's clock seem to be out of sync. For the statistics to be correct the different locust servers need to have synchronized clocks.")
             elif msg.type == "locustfile":
+                if msg.data["version"][0:4] == __version__[0:4]:
+                    logger.debug(
+                        f"A worker ({msg.node_id}) running a different patch version ({msg.data['version']}) connected, master version is {__version__}"
+                    )
+
                 logging.debug("Worker requested locust file")
-                assert self.environment.parsed_options
-                filename = self.environment.parsed_options.locustfile
+                assert self.environment.parsed_locustfiles
+                locustfile_options = self.environment.parsed_locustfiles
+                locustfile_list = [f.strip() for f in locustfile_options if not os.path.isdir(f)]
+
+                for locustfile_option in locustfile_options:
+                    if os.path.isdir(locustfile_option):
+                        locustfile_list.extend(get_abspaths_in(locustfile_option, extension=".py"))
+
                 try:
-                    with open(filename) as f:
-                        file_contents = f.read()
+                    locustfiles: list[str | dict[str, str]] = []
+
+                    for filename in locustfile_list:
+                        if is_url(filename):
+                            locustfiles.append(filename)
+                        else:
+                            with open(filename) as f:
+                                filename = os.path.basename(filename)
+                                file_contents = f.read()
+
+                            locustfiles.append({"filename": filename, "contents": file_contents})
                 except Exception as e:
-                    logger.error(
-                        f"--locustfile must be a full path to a single locustfile for file distribution to work {e}"
-                    )
+                    error_message = "locustfile must be a full path to a single locustfile, a comma-separated list of .py files, or a URL for file distribution to work"
+                    logger.error(f"{error_message} {e}")
                     self.send_message(
                         "locustfile",
                         client_id=client_id,
-                        data={
-                            "error": f"locustfile must be a full path to a single locustfile for file distribution to work (was '{filename}')"
-                        },
+                        data={"error": f"{error_message} (was '{filename}')"},
                     )
                 else:
-                    if getattr(self, "_old_file_contents", file_contents) != file_contents:
-                        logger.warning(
-                            "Locustfile contents changed on disk after first worker requested locustfile, sending new content. If you make any major changes (like changing User class names) you need to restart master."
-                        )
-                    self._old_file_contents = file_contents
                     self.send_message(
                         "locustfile",
                         client_id=client_id,
-                        data={"filename": os.path.basename(filename), "contents": file_contents},
+                        data={"locustfiles": locustfiles},
                     )
                 continue
             elif msg.type == "client_stopped":

locust/test/test_main.py

@@ -1721,11 +1721,6 @@ def t(self):
                 text=True,
             )
             gevent.sleep(2)
-            # modify the locustfile to trigger warning about file change when the second worker connects
-            with open(mocked.file_path, "w") as locustfile:
-                locustfile.write(LOCUSTFILE_CONTENT)
-                locustfile.write("\n# New comment\n")
-            gevent.sleep(2)
             proc_worker2 = subprocess.Popen(
                 [
                     "locust",
@@ -1742,7 +1737,6 @@ def t(self):
             stdout_worker2 = proc_worker2.communicate()[0]
 
             self.assertIn('All users spawned: {"User1": 1} (1 total users)', stdout)
-            self.assertIn("Locustfile contents changed on disk after first worker requested locustfile", stdout)
             self.assertIn("Shutting down (exit code 0)", stdout)
             self.assertNotIn("Traceback", stdout)
             self.assertNotIn("Traceback", stdout_worker)
@@ -1818,49 +1812,45 @@ def t(self):
             """
         )
         with mock_locustfile(content=LOCUSTFILE_CONTENT) as mocked:
-            with mock_locustfile() as mocked2:
-                proc = subprocess.Popen(
-                    [
-                        "locust",
-                        "-f",
-                        f"{mocked.file_path}, {mocked2.file_path}",
-                        "--headless",
-                        "--master",
-                        "-L",
-                        "debug",
-                    ],
-                    stderr=STDOUT,
-                    stdout=PIPE,
-                    text=True,
-                )
-                proc_worker = subprocess.Popen(
-                    [
-                        "locust",
-                        "-f",
-                        "-",
-                        "--worker",
-                    ],
-                    stderr=STDOUT,
-                    stdout=PIPE,
-                    text=True,
-                )
+            proc = subprocess.Popen(
+                [
+                    "locust",
+                    "-f",
+                    f"{mocked.file_path},{mocked.file_path}",
+                    "--headless",
+                    "--master",
+                    "--expect-workers",
+                    "1",
+                    "-t",
+                    "1s",
+                ],
+                stderr=STDOUT,
+                stdout=PIPE,
+                text=True,
+            )
+            gevent.sleep(2)
+            proc_worker = subprocess.Popen(
+                [
+                    "locust",
+                    "-f",
+                    "-",
+                    "--worker",
+                ],
+                stderr=STDOUT,
+                stdout=PIPE,
+                text=True,
+            )
 
-                try:
-                    stdout = proc_worker.communicate(timeout=5)[0]
-                    self.assertIn(
-                        "Got error from master: locustfile must be a full path to a single locustfile for file distribution to work",
-                        stdout,
-                    )
-                    proc.kill()
-                    master_stdout = proc.communicate()[0]
-                    self.assertIn(
-                        "--locustfile must be a full path to a single locustfile for file distribution", master_stdout
-                    )
-                except Exception:
-                    proc.kill()
-                    proc_worker.kill()
-                    stdout, worker_stderr = proc_worker.communicate()
-                    assert False, f"worker never finished: {stdout}"
+            stdout = proc.communicate()[0]
+            stdout_worker = proc_worker.communicate()[0]
+
+            self.assertIn('All users spawned: {"User1": 1} (1 total users)', stdout)
+            self.assertIn("Shutting down (exit code 0)", stdout)
+            self.assertNotIn("Traceback", stdout)
+            self.assertNotIn("Traceback", stdout_worker)
+
+            self.assertEqual(0, proc.returncode)
+            self.assertEqual(0, proc_worker.returncode)
 
     def test_json_schema(self):
         LOCUSTFILE_CONTENT = textwrap.dedent(

locust/util/directory.py

@@ -0,0 +1,12 @@
+import os
+
+
+def get_abspaths_in(path, extension=None):
+    return [
+        os.path.abspath(os.path.join(root, f))
+        for root, _dirs, fs in os.walk(path)
+        for f in fs
+        if os.path.isfile(os.path.join(root, f))
+        and (f.endswith(extension) or extension is None)
+        and not f.startswith("_")
+    ]

locust/util/url.py

@@ -0,0 +1,15 @@
+from urllib.parse import urlparse
+
+
+def is_url(url: str) -> bool:
+    """
+    Check if path is an url
+    """
+    try:
+        result = urlparse(url)
+        if result.scheme == "https" or result.scheme == "http":
+            return True
+        else:
+            return False
+    except ValueError:
+        return False

package.json

@@ -1,6 +1,5 @@
 {
   "name": "locust",
-  "version": "2.16.1",
   "private": true,
   "license": "MIT",
   "scripts": {