From 90967f9664363c15edf892471e6f8b993a17ad1a Mon Sep 17 00:00:00 2001
From: Michael Nester <mike.nester0@gmail.com>
Date: Wed, 24 Aug 2022 09:04:22 -0400
Subject: [PATCH] Added check to swarm for class picker. Uses pre selected user
 classes if test is already running

---
 locust/test/test_web.py | 53 +++++++++++++++++++++++++++++++++++++++++
 locust/web.py           | 20 +++++++++++-----
 2 files changed, 67 insertions(+), 6 deletions(-)

diff --git a/locust/test/test_web.py b/locust/test/test_web.py
index fa00727d88..80f9cd7e87 100644
--- a/locust/test/test_web.py
+++ b/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(
diff --git a/locust/web.py b/locust/web.py
index 48ad911fff..17eb11c22e 100644
--- a/locust/web.py
+++ b/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