From 34d8f57a2fea7766733ac51d2b9f3c16462d8bd9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tomasz=20D=C4=85dela?= <dadela.tomasz@gmail.com>
Date: Thu, 7 Nov 2024 00:20:12 +0100
Subject: [PATCH 1/2] fix: overwrite user weight by CL config before removing
 user classes with weight zero

---
 locust/env.py  |  4 ++++
 locust/main.py | 40 +++++++++++++++++++++-------------------
 2 files changed, 25 insertions(+), 19 deletions(-)

diff --git a/locust/env.py b/locust/env.py
index 9ba161c684..bec5e0e68f 100644
--- a/locust/env.py
+++ b/locust/env.py
@@ -37,6 +37,7 @@ def __init__(
         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,
+        override_user_weights: list[dict] | None = None,
         dispatcher_class: type[UsersDispatcher] = UsersDispatcher,
     ):
         self.runner: Runner | None = None
@@ -105,6 +106,9 @@ def __init__(
         self.worker_logs: dict[str, list[str]] = {}
         """Captured logs from all connected workers"""
 
+        for user_dict in override_user_weights:
+            self.update_user_class(user_dict)
+
         self._remove_user_classes_with_weight_zero()
         self._validate_user_class_name_uniqueness()
         self._validate_shape_class_instance()
diff --git a/locust/main.py b/locust/main.py
index 4031f2bb3b..ddb7d69229 100644
--- a/locust/main.py
+++ b/locust/main.py
@@ -69,6 +69,7 @@ def create_environment(
     available_user_classes=None,
     available_shape_classes=None,
     available_user_tasks=None,
+    override_user_weights=None,
 ):
     """
     Create an Environment instance from options
@@ -85,6 +86,7 @@ def create_environment(
         available_user_classes=available_user_classes,
         available_shape_classes=available_shape_classes,
         available_user_tasks=available_user_tasks,
+        override_user_weights=override_user_weights,
     )
 
 
@@ -355,21 +357,7 @@ def kill_workers(children):
 See https://github.com/locustio/locust/wiki/Installation#increasing-maximum-number-of-open-files-limit for more info."""
             )
 
-    # create locust Environment
-    locustfile_path = None if not locustfile else os.path.basename(locustfile)
-
-    environment = create_environment(
-        user_classes,
-        options,
-        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,
-    )
-
+    user_new_weights: list[dict] = []
     if options.config_users:
         for json_user_config in options.config_users:
             try:
@@ -387,12 +375,10 @@ def ensure_user_class_name(config):
                 if isinstance(user_config, list):
                     for config in user_config:
                         ensure_user_class_name(config)
-
-                        environment.update_user_class(config)
+                        user_new_weights.append(config)
                 else:
                     ensure_user_class_name(user_config)
-
-                    environment.update_user_class(user_config)
+                    user_new_weights.append(user_config)
             except json.decoder.JSONDecodeError as e:
                 logger.error(f"The --config-users argument must be in valid JSON string or file: {e}")
                 sys.exit(-1)
@@ -405,6 +391,22 @@ def ensure_user_class_name(config):
                 logger.exception(e)
                 sys.exit(-1)
 
+    # create locust Environment
+    locustfile_path = None if not locustfile else os.path.basename(locustfile)
+
+    environment = create_environment(
+        user_classes,
+        options,
+        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,
+        override_user_weights=user_new_weights,
+    )
+
     if (
         shape_class
         and not shape_class.use_common_options

From bbf8b013ec21e0a6052cf598b4afc8e39baaca0e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tomasz=20D=C4=85dela?= <dadela.tomasz@gmail.com>
Date: Thu, 7 Nov 2024 21:17:46 +0100
Subject: [PATCH 2/2] fix: check if override_user_weights is empty or None

---
 locust/env.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/locust/env.py b/locust/env.py
index 9fd75271b0..59c7467506 100644
--- a/locust/env.py
+++ b/locust/env.py
@@ -106,8 +106,9 @@ def __init__(
         self.worker_logs: dict[str, list[str]] = {}
         """Captured logs from all connected workers"""
 
-        for user_dict in override_user_weights:
-            self.update_user_class(user_dict)
+        if override_user_weights:
+            for user_dict in override_user_weights:
+                self.update_user_class(user_dict)
 
         self._remove_user_classes_with_weight_zero()
         self._validate_user_class_name_uniqueness()