Added check to swarm for class picker. Uses pre selected user classes…

… if test is already running

locust/test/test_web.py

@@ -391,6 +391,59 @@ def t(self):
         response = requests.get("http://127.0.0.1:%i/stop" % self.web_port)
         self.assertEqual(response.json()["message"], "Test stopped")
 
+    def test_swarm_uses_pre_selected_user_classes_when_empty_payload_and_test_is_already_running_with_class_picker(
+        self,
+    ):
+        # This test validates that the correct User Classes are used when editing a running test
+        class User1(User):
+            wait_time = constant(1)
+
+            @task
+            def t(self):
+                pass
+
+        class User2(User):
+            wait_time = constant(1)
+
+            @task
+            def t(self):
+                pass
+
+        self.environment.web_ui.userclass_picker_is_active = True
+        self.environment.available_user_classes = {"User1": User1, "User2": User2}
+        response = requests.post(
+            "http://127.0.0.1:%i/swarm" % self.web_port,
+            data={
+                "user_count": 5,
+                "spawn_rate": 5,
+                "host": "https://localhost",
+                "user_classes": ["User1"],
+            },
+        )
+
+        self.assertEqual(200, response.status_code)
+        self.assertEqual("https://localhost", response.json()["host"])
+        self.assertEqual(self.environment.host, "https://localhost")
+        self.assertListEqual(["User1"], response.json()["user_classes"])
+
+        # simulating edit running load test
+        response = requests.post(
+            "http://127.0.0.1:%i/swarm" % self.web_port,
+            data={
+                "user_count": 10,
+                "spawn_rate": 10,
+            },
+        )
+        self.assertEqual(200, response.status_code)
+        self.assertEqual("https://localhost", response.json()["host"])
+        self.assertEqual(self.environment.host, "https://localhost")
+        self.assertListEqual(["User1"], response.json()["user_classes"])
+
+        # stop
+        gevent.sleep(1)
+        response = requests.get("http://127.0.0.1:%i/stop" % self.web_port)
+        self.assertEqual(response.json()["message"], "Test stopped")
+
     def test_swarm_error_when_userclass_picker_is_active_but_no_available_userclasses(self):
         self.environment.web_ui.userclass_picker_is_active = True
         response = requests.post(

locust/web.py

@@ -1,7 +1,6 @@
 import csv
 import logging
 import os.path
-import sys
 from functools import wraps
 from html import escape
 from io import StringIO
@@ -14,10 +13,9 @@
 from flask import Flask, make_response, jsonify, render_template, request, send_file, Response
 from flask_basicauth import BasicAuth
 from gevent import pywsgi
-import locust
 
 from .exception import AuthCredentialsError
-from .runners import MasterRunner, STATE_MISSING
+from .runners import MasterRunner, STATE_RUNNING, STATE_MISSING
 from .log import greenlet_exception_logger
 from .stats import StatsCSVFileWriter, StatsErrorDict, sort_stats
 from . import stats as stats_module, __version__ as version, argument_parser
@@ -27,8 +25,6 @@
 from .util.rounding import proper_round
 from .html import get_html_report
 from flask_cors import CORS
-from .user.inspectuser import print_task_ratio, print_task_ratio_json
-from .util.timespan import parse_timespan
 
 if TYPE_CHECKING:
     from .env import Environment
@@ -164,7 +160,19 @@ def swarm() -> Response:
                             user_classes[user_class_name] = user_class_object
 
                 else:
-                    user_classes = self.environment.available_user_classes
+                    if self.environment.runner and self.environment.runner.state == STATE_RUNNING:
+                        # Test is already running
+                        # Using the user classes that have already been selected
+                        user_classes = {
+                            key: value
+                            for (key, value) in self.environment.available_user_classes.items()
+                            if value in self.environment.user_classes
+                        }
+                    else:
+                        # Starting test with no user class selection
+                        # Defaulting to using all available user classes
+                        user_classes = self.environment.available_user_classes
+
                 self._update_user_classes(user_classes)
 
                 # Updating ShapeClass if specified in WebUI Form