From f5410f0c734ab1c3c951cd2309c6d970a9ef19e8 Mon Sep 17 00:00:00 2001 From: Andrew Baldwin Date: Fri, 5 Apr 2024 09:07:48 +0200 Subject: [PATCH] Add env test --- locust/test/test_env.py | 40 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/locust/test/test_env.py b/locust/test/test_env.py index b545449bdb..9e468b2e4a 100644 --- 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}, + )