From ba7c49ea688eecc91334733224303c18e6459e9b Mon Sep 17 00:00:00 2001
From: Michael Nester <michaelnester@Michaels-MacBook-Pro-2.local>
Date: Mon, 30 Oct 2023 22:37:32 -0400
Subject: [PATCH 1/5] Userclass Picker styling fix

---
 locust/static/css/application.css | 17 +++++++++++++++++
 locust/templates/index.html       |  2 +-
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/locust/static/css/application.css b/locust/static/css/application.css
index 672a76b2ae..187eba4e51 100644
--- a/locust/static/css/application.css
+++ b/locust/static/css/application.css
@@ -470,5 +470,22 @@ body.ready .main {
   margin-right: 10px;
   color: #48a584;
 }
+.userClass {
+  border: none;
+  background: #fff;
+  height: 100px;
+  width: 340px;
+  font-size: 18px;
+  padding-left: 10px;
+}
+
+.shapeClass {
+  border: none;
+  background: #fff;
+  /* height: 100px; */
+  width: 340px;
+  font-size: 18px;
+  padding-left: 10px;
+}
 
 /*# sourceMappingURL=application.css.map */
diff --git a/locust/templates/index.html b/locust/templates/index.html
index 2abb59897d..0005887c1e 100644
--- a/locust/templates/index.html
+++ b/locust/templates/index.html
@@ -65,7 +65,7 @@ <h2>Start new load test</h2>
                         {% for user in available_user_classes %}
                             <option value="{{user}}" selected>{{user}}</option>
                         {% endfor %}
-                        </select><br>range
+                        </select><br>
                         <label for="available_shape_classes">ShapeClass </label>
                         <select name="shape_class" id="shape-classes" class="shapeClass">
                         {% for shape_class in available_shape_classes %}

From 5c1fc288106780cef27c1c7a6a6899eef39f60ed Mon Sep 17 00:00:00 2001
From: Michael Nester <michaelnester@Michaels-MacBook-Pro-2.local>
Date: Tue, 31 Oct 2023 10:12:16 -0400
Subject: [PATCH 2/5] fix for not loading all shape classes in class picker

---
 locust/main.py                 | 17 ++++++++++-------
 locust/util/load_locustfile.py | 13 +++++--------
 locust/web.py                  |  6 +++++-
 3 files changed, 20 insertions(+), 16 deletions(-)

diff --git a/locust/main.py b/locust/main.py
index 006ed58af4..d8a4643dcb 100644
--- a/locust/main.py
+++ b/locust/main.py
@@ -86,17 +86,20 @@ def main():
     user_classes: Dict[str, locust.User] = {}
     available_user_classes = {}
     available_shape_classes = {}
+    shape_class = None
     for _locustfile in locustfiles:
-        docstring, _user_classes, shape_class = load_locustfile(_locustfile)
+        docstring, _user_classes, shape_classes = load_locustfile(_locustfile)
 
         # Setting Available Shape Classes
-        if shape_class:
-            shape_class_name = type(shape_class).__name__
-            if shape_class_name in available_shape_classes.keys():
-                sys.stderr.write(f"Duplicate shape classes: {shape_class_name}\n")
-                sys.exit(1)
+        if shape_classes:
+            shape_class = shape_classes[0]
+            for shape_class in shape_classes:
+                shape_class_name = type(shape_class).__name__
+                if shape_class_name in available_shape_classes.keys():
+                    sys.stderr.write(f"Duplicate shape classes: {shape_class_name}\n")
+                    sys.exit(1)
 
-            available_shape_classes[shape_class_name] = shape_class
+                available_shape_classes[shape_class_name] = shape_class
 
         # Setting Available User Classes
         for key, value in _user_classes.items():
diff --git a/locust/util/load_locustfile.py b/locust/util/load_locustfile.py
index 6eae8d108b..06b772e488 100644
--- a/locust/util/load_locustfile.py
+++ b/locust/util/load_locustfile.py
@@ -2,7 +2,7 @@
 import inspect
 import os
 import sys
-from typing import Dict, Optional, Tuple
+from typing import Dict, List, Optional, Tuple
 from ..shape import LoadTestShape
 from ..user import User
 
@@ -21,7 +21,7 @@ def is_shape_class(item):
     return bool(inspect.isclass(item) and issubclass(item, LoadTestShape) and not getattr(item, "abstract", True))
 
 
-def load_locustfile(path) -> Tuple[Optional[str], Dict[str, User], Optional[LoadTestShape]]:
+def load_locustfile(path) -> Tuple[Optional[str], Dict[str, User], List[LoadTestShape]]:
     """
     Import given locustfile path and return (docstring, callables).
 
@@ -65,10 +65,7 @@ def load_locustfile(path) -> Tuple[Optional[str], Dict[str, User], Optional[Load
     user_classes = {name: value for name, value in vars(imported).items() if is_user_class(value)}
 
     # Find shape class, if any, return it
-    shape_classes = [value for name, value in vars(imported).items() if is_shape_class(value)]
-    if shape_classes:
-        shape_class = shape_classes[0]()
-    else:
-        shape_class = None
+    _shape_classes = [value for name, value in vars(imported).items() if is_shape_class(value)]
+    shape_classes = [shape_class() for shape_class in _shape_classes]
 
-    return imported.__doc__, user_classes, shape_class
+    return imported.__doc__, user_classes, shape_classes
diff --git a/locust/web.py b/locust/web.py
index f55f9fff17..6dd8511745 100644
--- a/locust/web.py
+++ b/locust/web.py
@@ -229,7 +229,11 @@ def swarm() -> Response:
             if environment.shape_class and environment.runner is not None:
                 environment.runner.start_shape()
                 return jsonify(
-                    {"success": True, "message": "Swarming started using shape class", "host": environment.host}
+                    {
+                        "success": True,
+                        "message": f"Swarming started using shape class '{type(environment.shape_class).__name__}'",
+                        "host": environment.host,
+                    }
                 )
 
             if self._swarm_greenlet is not None:

From 4aca017dc098d5e39568a669e5799481e502762a Mon Sep 17 00:00:00 2001
From: Michael Nester <michaelnester@Michaels-MacBook-Pro-2.local>
Date: Tue, 31 Oct 2023 10:12:45 -0400
Subject: [PATCH 3/5] formatting

---
 locust/util/load_locustfile.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/locust/util/load_locustfile.py b/locust/util/load_locustfile.py
index 06b772e488..f719786d2a 100644
--- a/locust/util/load_locustfile.py
+++ b/locust/util/load_locustfile.py
@@ -65,7 +65,6 @@ def load_locustfile(path) -> Tuple[Optional[str], Dict[str, User], List[LoadTest
     user_classes = {name: value for name, value in vars(imported).items() if is_user_class(value)}
 
     # Find shape class, if any, return it
-    _shape_classes = [value for name, value in vars(imported).items() if is_shape_class(value)]
-    shape_classes = [shape_class() for shape_class in _shape_classes]
+    shape_classes = [value() for name, value in vars(imported).items() if is_shape_class(value)]
 
     return imported.__doc__, user_classes, shape_classes

From bd8cfb9681f3229c1e2794851294e3e11419a9dd Mon Sep 17 00:00:00 2001
From: Michael Nester <michaelnester@Michaels-MacBook-Pro-2.local>
Date: Tue, 31 Oct 2023 10:40:59 -0400
Subject: [PATCH 4/5] test cases for shape class loading fix

---
 locust/test/test_load_locustfile.py | 40 +++++++++++++++++++++--------
 locust/test/test_web.py             | 18 +++++++++----
 2 files changed, 43 insertions(+), 15 deletions(-)

diff --git a/locust/test/test_load_locustfile.py b/locust/test/test_load_locustfile.py
index 79aa502c44..50c46163ed 100644
--- a/locust/test/test_load_locustfile.py
+++ b/locust/test/test_load_locustfile.py
@@ -37,25 +37,25 @@ class ThriftLocust(User):
 
     def test_load_locust_file_from_absolute_path(self):
         with mock_locustfile() as mocked:
-            docstring, user_classes, shape_class = main.load_locustfile(mocked.file_path)
+            docstring, user_classes, shape_classes = main.load_locustfile(mocked.file_path)
             self.assertIn("UserSubclass", user_classes)
             self.assertNotIn("NotUserSubclass", user_classes)
             self.assertNotIn("LoadTestShape", user_classes)
-            self.assertIsNone(shape_class)
+            self.assertEqual(shape_classes, [])
 
     def test_load_locust_file_from_relative_path(self):
         with mock_locustfile() as mocked:
-            docstring, user_classes, shape_class = main.load_locustfile(
+            docstring, user_classes, shape_classes = main.load_locustfile(
                 os.path.join(os.path.relpath(mocked.directory, os.getcwd()), mocked.filename)
             )
 
     def test_load_locust_file_with_a_dot_in_filename(self):
         with mock_locustfile(filename_prefix="mocked.locust.file") as mocked:
-            docstring, user_classes, shape_class = main.load_locustfile(mocked.file_path)
+            docstring, user_classes, shape_classes = main.load_locustfile(mocked.file_path)
 
     def test_return_docstring_and_user_classes(self):
         with mock_locustfile() as mocked:
-            docstring, user_classes, shape_class = main.load_locustfile(mocked.file_path)
+            docstring, user_classes, shape_classes = main.load_locustfile(mocked.file_path)
             self.assertEqual("This is a mock locust file for unit testing", docstring)
             self.assertIn("UserSubclass", user_classes)
             self.assertNotIn("NotUserSubclass", user_classes)
@@ -69,12 +69,32 @@ def tick(self):
                     return None
         """
         )
+        with mock_locustfile(content=content) as mocked:
+            docstring, user_classes, shape_classes = main.load_locustfile(mocked.file_path)
+            self.assertEqual("This is a mock locust file for unit testing", docstring)
+            self.assertIn("UserSubclass", user_classes)
+            self.assertNotIn("NotUserSubclass", user_classes)
+            self.assertEqual(shape_classes[0].__class__.__name__, "LoadTestShape")
+
+    def test_with_multiple_shape_classes(self):
+        content = MOCK_LOCUSTFILE_CONTENT + textwrap.dedent(
+            """\
+        class LoadTestShape1(LoadTestShape):
+            def tick(self):
+                pass
+
+        class LoadTestShape2(LoadTestShape):
+            def tick(self):
+                pass
+        """
+        )
         with mock_locustfile(content=content) as mocked:
             docstring, user_classes, shape_class = main.load_locustfile(mocked.file_path)
             self.assertEqual("This is a mock locust file for unit testing", docstring)
             self.assertIn("UserSubclass", user_classes)
             self.assertNotIn("NotUserSubclass", user_classes)
-            self.assertEqual(shape_class.__class__.__name__, "LoadTestShape")
+            self.assertEqual(shape_class[0].__class__.__name__, "LoadTestShape1")
+            self.assertEqual(shape_class[1].__class__.__name__, "LoadTestShape2")
 
     def test_with_abstract_shape_class(self):
         content = MOCK_LOCUSTFILE_CONTENT + textwrap.dedent(
@@ -92,10 +112,10 @@ class UserLoadTestShape(UserBaseLoadTestShape):
         )
 
         with mock_locustfile(content=content) as mocked:
-            _, user_classes, shape_class = main.load_locustfile(mocked.file_path)
+            _, user_classes, shape_classes = main.load_locustfile(mocked.file_path)
             self.assertNotIn("UserBaseLoadTestShape", user_classes)
             self.assertNotIn("UserLoadTestShape", user_classes)
-            self.assertEqual(shape_class.__class__.__name__, "UserLoadTestShape")
+            self.assertEqual(shape_classes[0].__class__.__name__, "UserLoadTestShape")
 
     def test_with_not_imported_shape_class(self):
         content = MOCK_LOCUSTFILE_CONTENT + textwrap.dedent(
@@ -107,9 +127,9 @@ def tick(self):
         )
 
         with mock_locustfile(content=content) as mocked:
-            _, user_classes, shape_class = main.load_locustfile(mocked.file_path)
+            _, user_classes, shape_classes = main.load_locustfile(mocked.file_path)
             self.assertNotIn("UserLoadTestShape", user_classes)
-            self.assertEqual(shape_class.__class__.__name__, "UserLoadTestShape")
+            self.assertEqual(shape_classes[0].__class__.__name__, "UserLoadTestShape")
 
     def test_create_environment(self):
         options = parse_options(
diff --git a/locust/test/test_web.py b/locust/test/test_web.py
index 545e1c27d4..ff4d74f09b 100644
--- a/locust/test/test_web.py
+++ b/locust/test/test_web.py
@@ -563,7 +563,15 @@ class User2(User):
             def t(self):
                 pass
 
-        class TestShape(LoadTestShape):
+        class TestShape1(LoadTestShape):
+            def tick(self):
+                run_time = self.get_run_time()
+                if run_time < 10:
+                    return 4, 4
+                else:
+                    return None
+
+        class TestShape2(LoadTestShape):
             def tick(self):
                 run_time = self.get_run_time()
                 if run_time < 10:
@@ -573,8 +581,8 @@ def tick(self):
 
         self.environment.web_ui.userclass_picker_is_active = True
         self.environment.available_user_classes = {"User1": User1, "User2": User2}
-        self.environment.available_shape_classes = {"TestShape": TestShape()}
-        self.environment.shape_class = None
+        self.environment.available_shape_classes = {"TestShape1": TestShape1(), "TestShape2": TestShape2()}
+        self.environment.shape_class = TestShape1()
 
         response = requests.post(
             "http://127.0.0.1:%i/swarm" % self.web_port,
@@ -583,14 +591,14 @@ def tick(self):
                 "spawn_rate": 5,
                 "host": "https://localhost",
                 "user_classes": "User1",
-                "shape_class": "TestShape",
+                "shape_class": "TestShape2",
             },
         )
 
         self.assertEqual(200, response.status_code)
         self.assertEqual("https://localhost", response.json()["host"])
         self.assertEqual(self.environment.host, "https://localhost")
-        assert isinstance(self.environment.shape_class, TestShape)
+        assert isinstance(self.environment.shape_class, TestShape2)
 
         # stop
         gevent.sleep(1)

From 0ff132c124aa5217fd7aeaf91dab43eec585c14d Mon Sep 17 00:00:00 2001
From: Michael Nester <michaelnester@Michaels-MacBook-Pro-2.local>
Date: Tue, 31 Oct 2023 10:44:33 -0400
Subject: [PATCH 5/5] typo fix

---
 locust/test/test_load_locustfile.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/locust/test/test_load_locustfile.py b/locust/test/test_load_locustfile.py
index 50c46163ed..376e2e6b00 100644
--- a/locust/test/test_load_locustfile.py
+++ b/locust/test/test_load_locustfile.py
@@ -89,12 +89,12 @@ def tick(self):
         """
         )
         with mock_locustfile(content=content) as mocked:
-            docstring, user_classes, shape_class = main.load_locustfile(mocked.file_path)
+            docstring, user_classes, shape_classes = main.load_locustfile(mocked.file_path)
             self.assertEqual("This is a mock locust file for unit testing", docstring)
             self.assertIn("UserSubclass", user_classes)
             self.assertNotIn("NotUserSubclass", user_classes)
-            self.assertEqual(shape_class[0].__class__.__name__, "LoadTestShape1")
-            self.assertEqual(shape_class[1].__class__.__name__, "LoadTestShape2")
+            self.assertEqual(shape_classes[0].__class__.__name__, "LoadTestShape1")
+            self.assertEqual(shape_classes[1].__class__.__name__, "LoadTestShape2")
 
     def test_with_abstract_shape_class(self):
         content = MOCK_LOCUSTFILE_CONTENT + textwrap.dedent(