Add env test

locustio · Apr 5, 2024 · f5410f0 · f5410f0
1 parent f57e004
commit f5410f0
Showing 1 changed file with 40 additions and 0 deletions.
diff --git a/locust/test/test_env.py b/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},
+        )