Update shape class' runner when Web UI picker is used

locust/test/test_web.py

@@ -8,6 +8,7 @@
 import logging
 from io import StringIO
 from tempfile import NamedTemporaryFile, TemporaryDirectory
+from typing import List, Optional, Tuple, Type
 
 import gevent
 import requests
@@ -656,6 +657,41 @@ def tick(self):
         response = requests.get("http://127.0.0.1:%i/stop" % self.web_port)
         self.assertEqual(response.json()["message"], "Test stopped")
 
+    def test_swarm_shape_class_is_updated_when_userclass_picker_is_active(self):
+        class User1(User):
+            pass
+
+        class TestShape(LoadTestShape):
+            def tick(self):
+                pass
+
+        test_shape_instance = TestShape()
+
+        self.environment.web_ui.userclass_picker_is_active = True
+        self.environment.available_user_classes = {"User1": User1}
+        self.environment.available_shape_classes = {"TestShape": test_shape_instance}
+        self.environment.shape_class = None
+
+        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",
+                "shape_class": "TestShape",
+            },
+        )
+
+        self.assertEqual(200, response.status_code)
+        self.assertEqual(test_shape_instance, self.environment.shape_class)
+        self.assertIsNotNone(test_shape_instance.runner)
+
+        # 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_userclass_shapeclass_ignored_when_userclass_picker_is_inactive(self):
         class User1(User):
             wait_time = constant(1)

locust/web.py

@@ -629,6 +629,7 @@ def update_template_args(self):
     def _update_shape_class(self, shape_class_name):
         if shape_class_name:
             shape_class = self.environment.available_shape_classes[shape_class_name]
+            shape_class.runner = self.environment.runner
         else:
             shape_class = None