Merge pull request #2666 from andrewbaldwin44/bugfix/2657

Update User Classes in Distributed Mode

Author: cyberw

Diff for: locust/env.py

@@ -198,6 +198,9 @@ def create_web_ui(
         return self.web_ui
 
     def update_user_class(self, user_settings):
+        if isinstance(self.runner, MasterRunner):
+            self.runner.send_message("update_user_class", user_settings)
+
         user_class_name = user_settings.get("user_class_name")
         user_class = self.available_user_classes[user_class_name]
         user_tasks = self.available_user_tasks[user_class_name]

Diff for: locust/runners.py

@@ -1395,6 +1395,8 @@ def worker(self) -> NoReturn:
                 self.reset_connection()
             elif msg.type == "heartbeat":
                 self.last_heartbeat_timestamp = time.time()
+            elif msg.type == "update_user_class":
+                self.environment.update_user_class(msg.data)
             elif msg.type in self.custom_messages:
                 logger.debug("Received %s message from master" % msg.type)
                 listener, concurrent = self.custom_messages[msg.type]

Diff for: locust/test/test_env.py

@@ -241,3 +241,43 @@ def my_task(self):
             environment.available_user_classes["User1"].json(),
             {"host": "http://localhost", "tasks": ["my_task_2"], "fixed_count": 0, "weight": 1},
         )
+
+    def test_distributed_update_user_class(self):
+        class MyUser1(User):
+            @task
+            def my_task(self):
+                pass
+
+            @task
+            def my_task_2(self):
+                pass
+
+        class MyUser2(User):
+            @task
+            def my_task(self):
+                pass
+
+        master_env = Environment(
+            user_classes=[MyUser1, MyUser2],
+            available_user_classes={"User1": MyUser1, "User2": MyUser2},
+            available_user_tasks={"User1": MyUser1.tasks, "User2": MyUser2.tasks},
+        )
+        master = master_env.create_master_runner("*", 0)
+
+        worker_env = Environment(
+            user_classes=[MyUser1, MyUser2],
+            available_user_classes={"User1": MyUser1, "User2": MyUser2},
+            available_user_tasks={"User1": MyUser1.tasks, "User2": MyUser2.tasks},
+        )
+        worker = worker_env.create_worker_runner("127.0.0.1", master.server.port)
+
+        master_env.update_user_class({"user_class_name": "User1", "host": "http://localhost", "tasks": ["my_task_2"]})
+
+        self.assertEqual(
+            master_env.available_user_classes["User1"].json(),
+            {"host": "http://localhost", "tasks": ["my_task_2"], "fixed_count": 0, "weight": 1},
+        )
+        self.assertEqual(
+            worker_env.available_user_classes["User1"].json(),
+            {"host": "http://localhost", "tasks": ["my_task_2"], "fixed_count": 0, "weight": 1},
+        )

Diff for: locust/test/test_web.py

@@ -1351,3 +1351,32 @@ def test_html_stats_report(self):
         self.assertTrue(d("#root"))
         self.assertIn('"locustfile": "locust.py"', str(d))
         self.assertIn('"host": "http://localhost"', str(d))
+
+    def test_update_user(self):
+        class MyUser1(User):
+            @task
+            def my_task(self):
+                pass
+
+            @task
+            def my_task_2(self):
+                pass
+
+        class MyUser2(User):
+            @task
+            def my_task(self):
+                pass
+
+        self.environment.user_classes = [MyUser1, MyUser2]
+        self.environment.available_user_classes = {"User1": MyUser1, "User2": MyUser2}
+        self.environment.available_user_tasks = {"User1": MyUser1.tasks, "User2": MyUser2.tasks}
+
+        requests.post(
+            "http://127.0.0.1:%i/user" % self.web_port,
+            json={"user_class_name": "User1", "host": "http://localhost", "tasks": ["my_task_2"]},
+        )
+
+        self.assertEqual(
+            self.environment.available_user_classes["User1"].json(),
+            {"host": "http://localhost", "tasks": ["my_task_2"], "fixed_count": 0, "weight": 1},
+        )