From 3dbd3fe16b9b09c36002161e894533f8eaaa410f Mon Sep 17 00:00:00 2001
From: Lars Holmberg <Lars.Holmberg@svenskaspel.se>
Date: Mon, 19 Feb 2024 18:56:36 +0100
Subject: [PATCH] Use built-in locustfile distribution instead of rsync. Change
 plugins-upload to default to off (replacing --skip-plugins with
 --upload-plugins)

---
 locust_swarm/swarm.py | 35 ++++++++++++++++++++++-------------
 setup.py              | 10 +---------
 2 files changed, 23 insertions(+), 22 deletions(-)

diff --git a/locust_swarm/swarm.py b/locust_swarm/swarm.py
index 727e4c8..f080125 100644
--- a/locust_swarm/swarm.py
+++ b/locust_swarm/swarm.py
@@ -143,12 +143,18 @@
     default=[],
     help="A list of extra files or directories to upload. Space-separated, e.g. --extra-files testdata.csv *.py my-directory/",
 )
-parser.add_argument(  # secret parameter to optimize by not uploading locust-plugins every time
+parser.add_argument(
     "--skip-plugins",
     action="store_true",
     default=False,
     help=configargparse.SUPPRESS,
 )
+parser.add_argument(
+    "--upload-plugins",
+    action="store_true",
+    default=False,
+    help="Upload locust-plugins to load gens (useful if you are developing locust-plugins)",
+)
 
 parser.add_argument(
     "--version",
@@ -248,8 +254,8 @@ def cleanup(server_list):
 
 
 def upload(server):
-    files = [args.locustfile or "locustfile.py"] + args.extra_files
-    if not args.skip_plugins:
+    files = args.extra_files.copy()
+    if args.upload_plugins:
         try:
             files.append(os.path.dirname(svs_locust.__file__))
         except NameError:
@@ -260,7 +266,12 @@ def upload(server):
 
             files.append(os.path.dirname(locust_plugins.__file__))
         except ModuleNotFoundError:
-            logging.debug("locust-plugins wasnt installed, not adding it to files to upload")
+            logging.error("locust-plugins wasnt installed")
+            sys.exit(1)
+
+    if not files:
+        return
+
     if len(files) > 1:
         filestr = "{" + ",".join(files) + "}"
     else:
@@ -272,7 +283,7 @@ def upload(server):
         check_output(f"rsync -qrtl --exclude __pycache__ --exclude .mypy_cache {filestr} {server}:")
 
 
-def start_worker_process(server, port, locustfile_filename):
+def start_worker_process(server, port):
     upload(server)
 
     if args.selenium:
@@ -344,7 +355,7 @@ def start_worker_process(server, port, locustfile_filename):
             "--expect-workers-max-wait",
             "30",
             "-f",
-            locustfile_filename,
+            "-",
             *ensure_remote_kill,
             "'",
         ]
@@ -375,17 +386,15 @@ def main():
 
     locustfile = args.locustfile or "locustfile.py"
 
-    if "/" in locustfile:
-        parser.error(
-            "Locustfile (-f) must be a file in the current directory (I'm lazy and havent fixed support for this yet)"
-        )
-
-    locustfile_filename = os.path.split(locustfile)[1]
     port = int(args.port)
     if args.processes_per_loadgen:
         parser.error(
             f"--processes-per-loadgen has been removed in favour of locusts native --processes parameter (you had it set to {args.processes_per_loadgen})"
         )
+    if args.skip_plugins:
+        parser.error(
+            "--skip-plugins has been removed, the default is now NOT to upload plugins (but you can enable it with --upload-plugins)"
+        )
     worker_process_count = args.processes * args.loadgens
     loadgen_list = args.loadgen_list.split(",")
     if args.loadgens < 0:
@@ -503,7 +512,7 @@ def get_available_servers_and_lock_them():
     for server in server_list:
         # fail early if master has already terminated
         check_proc_running(master_proc)
-        worker_procs.extend(start_worker_process(server, port, locustfile_filename))
+        worker_procs.extend(start_worker_process(server, port))
 
     # check that worker procs didnt immediately terminate for some reason (like invalid parameters)
     try:
diff --git a/setup.py b/setup.py
index d99b310..fc18752 100644
--- a/setup.py
+++ b/setup.py
@@ -2,8 +2,6 @@
 # -*- coding: utf-8 -*-
 import os
 import sys
-
-import importlib.util
 from setuptools import find_packages, setup
 from setuptools.command.install import install
 from setuptools.command.egg_info import egg_info
@@ -31,12 +29,6 @@ def run(self):
         install_check(self, egg_info)
 
 
-requirement_list = ["locust>=2.19.0"]
-# if locust-plugins IS installed, then require a version known to work with this version of swarm.
-spec = importlib.util.find_spec("locust_plugins")
-if spec is not None:
-    requirement_list.append("locust-plugins>=2.7.0")
-
 setup(
     name="locust-swarm",
     description="Load test + test data distribution & launching tool for Locust",
@@ -65,7 +57,7 @@ def run(self):
     packages=find_packages(),
     include_package_data=True,
     zip_safe=False,
-    install_requires=requirement_list,
+    install_requires=["locust>=2.23.0"],
     entry_points={
         "console_scripts": ["swarm = locust_swarm.swarm:main"],
     },