From d3715324103959b87998ad3bafb77a217f2eb7d7 Mon Sep 17 00:00:00 2001 From: Shekar Ramaswamy Date: Thu, 5 Aug 2021 16:44:47 -0700 Subject: [PATCH 1/7] v1 working, need tests and cleaning --- locust/argument_parser.py | 7 +++++++ locust/main.py | 29 +++++++++++++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/locust/argument_parser.py b/locust/argument_parser.py index d7f617a7e9..befbc8c5e4 100644 --- a/locust/argument_parser.py +++ b/locust/argument_parser.py @@ -417,6 +417,13 @@ def setup_parser_arguments(parser): help="Number of seconds to wait for a simulated user to complete any executing task before exiting. Default is to terminate immediately. This parameter only needs to be specified for the master process when running Locust distributed.", env_var="LOCUST_STOP_TIMEOUT", ) + other_group.add_argument( + "--use-uniform-task-weights", + action="store_true", + default=False, + dest="use_uniform_task_weights", + help="Use uniformly distributed task weights, overriding the weights specified in the locustfile.", + ) user_classes_group = parser.add_argument_group("User classes") user_classes_group.add_argument( diff --git a/locust/main.py b/locust/main.py index 01fea3dd94..0327ac6e12 100644 --- a/locust/main.py +++ b/locust/main.py @@ -19,6 +19,7 @@ from .stats import print_error_report, print_percentile_stats, print_stats, stats_printer, stats_history from .stats import StatsCSV, StatsCSVFileWriter from .user import User +from .user.task import TaskSet from .user.inspectuser import get_task_ratio_dict, print_task_ratio from .util.timespan import parse_timespan from .exception import AuthCredentialsError @@ -298,6 +299,34 @@ def main(): else: web_ui = None + def assign_uniform_task_weights(environment, **kwargs): + for u in environment.user_classes: + u.weight = 1 + user_tasks = [] + tasks_frontier = u.tasks + # print("starting length for user tasks " + str(len(u.tasks))) + while len(tasks_frontier) != 0: + t = tasks_frontier.pop() + if hasattr(t, "tasks") and t.tasks: + # print("adding " + str(len(t.tasks))) + # print(t.tasks) + # print('---') + tasks_frontier.extend(t.tasks) + elif isinstance(t, Callable): + # print("examiming task") + # print(t) + # print('---') + if t not in user_tasks: + user_tasks.append(t) + else: + logger.error("Unrecognized type in user tasks") + # print("final num tasks") + # print(len(user_asks)) + u.tasks = user_tasks + + if options.use_uniform_task_weights: + environment.events.init.add_listener(assign_uniform_task_weights)t + # Fire locust init event which can be used by end-users' code to run setup code that # need access to the Environment, Runner or WebUI. environment.events.init.fire(environment=environment, runner=runner, web_ui=web_ui) From 41f0b09258bfbf39a5b1512600ae27dc338a57c1 Mon Sep 17 00:00:00 2001 From: Shekar Ramaswamy Date: Thu, 5 Aug 2021 17:43:41 -0700 Subject: [PATCH 2/7] Updates --- locust/argument_parser.py | 4 ++-- locust/main.py | 7 ++++--- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/locust/argument_parser.py b/locust/argument_parser.py index befbc8c5e4..5c12b38cea 100644 --- a/locust/argument_parser.py +++ b/locust/argument_parser.py @@ -418,10 +418,10 @@ def setup_parser_arguments(parser): env_var="LOCUST_STOP_TIMEOUT", ) other_group.add_argument( - "--use-uniform-task-weights", + "--same-task-weights", action="store_true", default=False, - dest="use_uniform_task_weights", + dest="same_task_weights", help="Use uniformly distributed task weights, overriding the weights specified in the locustfile.", ) diff --git a/locust/main.py b/locust/main.py index 0327ac6e12..b1e6f5bb59 100644 --- a/locust/main.py +++ b/locust/main.py @@ -6,6 +6,7 @@ import socket import sys import time +from typing import Callable import gevent @@ -299,7 +300,7 @@ def main(): else: web_ui = None - def assign_uniform_task_weights(environment, **kwargs): + def assign_same_task_weights(environment, **kwargs): for u in environment.user_classes: u.weight = 1 user_tasks = [] @@ -324,8 +325,8 @@ def assign_uniform_task_weights(environment, **kwargs): # print(len(user_asks)) u.tasks = user_tasks - if options.use_uniform_task_weights: - environment.events.init.add_listener(assign_uniform_task_weights)t + if options.same_task_weights: + environment.events.init.add_listener(assign_same_task_weights) # Fire locust init event which can be used by end-users' code to run setup code that # need access to the Environment, Runner or WebUI. From 16ea9bb3b67e6ee2542f67c9c955d1053727cba0 Mon Sep 17 00:00:00 2001 From: Shekar Ramaswamy Date: Fri, 6 Aug 2021 10:23:11 -0700 Subject: [PATCH 3/7] added tests --- locust/env.py | 23 ++++++++++- locust/main.py | 26 +------------ locust/test/test_env.py | 85 +++++++++++++++++++++++++++++++++++++++++ 3 files changed, 108 insertions(+), 26 deletions(-) diff --git a/locust/env.py b/locust/env.py index 05b9e64e13..814a1947d0 100644 --- a/locust/env.py +++ b/locust/env.py @@ -1,5 +1,6 @@ from operator import methodcaller from typing import ( + Callable, Dict, List, Type, @@ -12,7 +13,7 @@ from .runners import Runner, LocalRunner, MasterRunner, WorkerRunner from .web import WebUI from .user import User -from .user.task import filter_tasks_by_tags +from .user.task import TaskSet, filter_tasks_by_tags from .shape import LoadTestShape @@ -215,6 +216,26 @@ def _filter_tasks_by_tags(self): for user_class in self.user_classes: filter_tasks_by_tags(user_class, self.tags, self.exclude_tags) + def assign_same_task_weights(self): + """ + Update the user classes such that each user runs their specified tasks with equal + probability. + """ + for u in self.user_classes: + u.weight = 1 + user_tasks = [] + tasks_frontier = u.tasks + while len(tasks_frontier) != 0: + t = tasks_frontier.pop() + if hasattr(t, "tasks") and t.tasks: + tasks_frontier.extend(t.tasks) + elif isinstance(t, Callable): + if t not in user_tasks: + user_tasks.append(t) + else: + raise ValueError("Unrecognized task type in user") + u.tasks = user_tasks + @property def user_classes_by_name(self) -> Dict[str, Type[User]]: return {u.__name__: u for u in self.user_classes} diff --git a/locust/main.py b/locust/main.py index b1e6f5bb59..8f43911849 100644 --- a/locust/main.py +++ b/locust/main.py @@ -6,7 +6,6 @@ import socket import sys import time -from typing import Callable import gevent @@ -20,7 +19,6 @@ from .stats import print_error_report, print_percentile_stats, print_stats, stats_printer, stats_history from .stats import StatsCSV, StatsCSVFileWriter from .user import User -from .user.task import TaskSet from .user.inspectuser import get_task_ratio_dict, print_task_ratio from .util.timespan import parse_timespan from .exception import AuthCredentialsError @@ -301,29 +299,7 @@ def main(): web_ui = None def assign_same_task_weights(environment, **kwargs): - for u in environment.user_classes: - u.weight = 1 - user_tasks = [] - tasks_frontier = u.tasks - # print("starting length for user tasks " + str(len(u.tasks))) - while len(tasks_frontier) != 0: - t = tasks_frontier.pop() - if hasattr(t, "tasks") and t.tasks: - # print("adding " + str(len(t.tasks))) - # print(t.tasks) - # print('---') - tasks_frontier.extend(t.tasks) - elif isinstance(t, Callable): - # print("examiming task") - # print(t) - # print('---') - if t not in user_tasks: - user_tasks.append(t) - else: - logger.error("Unrecognized type in user tasks") - # print("final num tasks") - # print(len(user_asks)) - u.tasks = user_tasks + environment.assign_same_task_weights() if options.same_task_weights: environment.events.init.add_listener(assign_same_task_weights) diff --git a/locust/test/test_env.py b/locust/test/test_env.py index 099c7ee6b1..577c6f93bf 100644 --- a/locust/test/test_env.py +++ b/locust/test/test_env.py @@ -6,6 +6,7 @@ User, task, ) +from locust.user.task import TaskSet from .testcases import LocustTestCase from .fake_module1_for_env_test import MyUserWithSameName as MyUserWithSameName1 from .fake_module2_for_env_test import MyUserWithSameName as MyUserWithSameName2 @@ -39,3 +40,87 @@ def test_user_classes_with_same_name_is_error(self): e.exception.args[0], "The following user classes have the same class name: locust.test.fake_module1_for_env_test.MyUserWithSameName, locust.test.fake_module2_for_env_test.MyUserWithSameName", ) + + def test_assign_same_task_weights(self): + def verify_tasks(u, target_tasks): + self.assertEqual(len(u.tasks), len(target_tasks)) + tasks = [t.__name__ for t in u.tasks] + self.assertEqual(len(tasks), len(set(tasks))) + self.assertEqual(set(tasks), set(target_tasks)) + + # Base case + class MyUser1(User): + wait_time = constant(0) + + @task(4) + def my_task(self): + pass + + @task(1) + def my_task_2(self): + pass + + environment = Environment(user_classes=[MyUser1]) + environment.assign_same_task_weights() + u = environment.user_classes[0] + verify_tasks(u, ["my_task", "my_task_2"]) + + # Testing nested task sets + class MyUser2(User): + @task + class TopLevelTaskSet(TaskSet): + @task + class IndexTaskSet(TaskSet): + @task(10) + def index(self): + self.client.get("/") + + @task + def stop(self): + self.client.get("/hi") + + @task(2) + def stats(self): + self.client.get("/stats/requests") + + environment = Environment(user_classes=[MyUser2]) + environment.assign_same_task_weights() + u = environment.user_classes[0] + verify_tasks(u, ["index", "stop", "stats"]) + + # Testing task assignment via instance variable + def outside_task(): + pass + def outside_task_2(): + pass + class SingleTaskSet(TaskSet): + tasks = [outside_task, outside_task, outside_task_2] + class MyUser3(User): + tasks = [SingleTaskSet, outside_task] + + environment = Environment(user_classes=[MyUser3]) + environment.assign_same_task_weights() + u = environment.user_classes[0] + verify_tasks(u, ["outside_task", "outside_task_2"]) + + # Testing task assignment via dict + class DictTaskSet(TaskSet): + def dict_task_1(): + pass + def dict_task_2(): + pass + def dict_task_3(): + pass + tasks = { + dict_task_1: 5, + dict_task_2: 3, + dict_task_3: 1, + } + class MyUser4(User): + tasks = [DictTaskSet, SingleTaskSet, SingleTaskSet] + + environment = Environment(user_classes=[MyUser4]) + environment.assign_same_task_weights() + u = environment.user_classes[0] + verify_tasks(u, ["outside_task", "outside_task_2", "dict_task_1", "dict_task_2", + "dict_task_3"]) From 48c35f37b2c9ff9529820b78eb5358cca28c51e1 Mon Sep 17 00:00:00 2001 From: Shekar Ramaswamy Date: Fri, 6 Aug 2021 10:32:34 -0700 Subject: [PATCH 4/7] Lint --- locust/test/test_env.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/locust/test/test_env.py b/locust/test/test_env.py index 577c6f93bf..633fb5cfa1 100644 --- a/locust/test/test_env.py +++ b/locust/test/test_env.py @@ -77,24 +77,27 @@ def index(self): @task def stop(self): - self.client.get("/hi") + self.client.get("/hi") @task(2) def stats(self): self.client.get("/stats/requests") - + environment = Environment(user_classes=[MyUser2]) environment.assign_same_task_weights() u = environment.user_classes[0] verify_tasks(u, ["index", "stop", "stats"]) - + # Testing task assignment via instance variable def outside_task(): pass + def outside_task_2(): pass + class SingleTaskSet(TaskSet): tasks = [outside_task, outside_task, outside_task_2] + class MyUser3(User): tasks = [SingleTaskSet, outside_task] @@ -107,20 +110,23 @@ class MyUser3(User): class DictTaskSet(TaskSet): def dict_task_1(): pass + def dict_task_2(): pass + def dict_task_3(): pass + tasks = { dict_task_1: 5, dict_task_2: 3, dict_task_3: 1, } + class MyUser4(User): tasks = [DictTaskSet, SingleTaskSet, SingleTaskSet] environment = Environment(user_classes=[MyUser4]) environment.assign_same_task_weights() u = environment.user_classes[0] - verify_tasks(u, ["outside_task", "outside_task_2", "dict_task_1", "dict_task_2", - "dict_task_3"]) + verify_tasks(u, ["outside_task", "outside_task_2", "dict_task_1", "dict_task_2", "dict_task_3"]) From 35ac94bfa3853ab2723719d3de54a82c03e5451a Mon Sep 17 00:00:00 2001 From: Shekar Ramaswamy Date: Fri, 6 Aug 2021 16:02:19 -0700 Subject: [PATCH 5/7] Updates --- locust/argument_parser.py | 4 ++-- locust/env.py | 2 +- locust/main.py | 8 ++++---- locust/test/test_env.py | 23 ++++++++++++++++++----- 4 files changed, 25 insertions(+), 12 deletions(-) diff --git a/locust/argument_parser.py b/locust/argument_parser.py index 5c12b38cea..1503811425 100644 --- a/locust/argument_parser.py +++ b/locust/argument_parser.py @@ -418,10 +418,10 @@ def setup_parser_arguments(parser): env_var="LOCUST_STOP_TIMEOUT", ) other_group.add_argument( - "--same-task-weights", + "--same-weights", action="store_true", default=False, - dest="same_task_weights", + dest="same_weights", help="Use uniformly distributed task weights, overriding the weights specified in the locustfile.", ) diff --git a/locust/env.py b/locust/env.py index 814a1947d0..101264bcc0 100644 --- a/locust/env.py +++ b/locust/env.py @@ -216,7 +216,7 @@ def _filter_tasks_by_tags(self): for user_class in self.user_classes: filter_tasks_by_tags(user_class, self.tags, self.exclude_tags) - def assign_same_task_weights(self): + def assign_uniform_weights(self): """ Update the user classes such that each user runs their specified tasks with equal probability. diff --git a/locust/main.py b/locust/main.py index 8f43911849..e609a4a953 100644 --- a/locust/main.py +++ b/locust/main.py @@ -298,11 +298,11 @@ def main(): else: web_ui = None - def assign_same_task_weights(environment, **kwargs): - environment.assign_same_task_weights() + def assign_uniform_weights(environment, **kwargs): + environment.assign_uniform_weights() - if options.same_task_weights: - environment.events.init.add_listener(assign_same_task_weights) + if options.same_weights: + environment.events.init.add_listener(assign_uniform_weights) # Fire locust init event which can be used by end-users' code to run setup code that # need access to the Environment, Runner or WebUI. diff --git a/locust/test/test_env.py b/locust/test/test_env.py index 633fb5cfa1..84558e5960 100644 --- a/locust/test/test_env.py +++ b/locust/test/test_env.py @@ -41,7 +41,7 @@ def test_user_classes_with_same_name_is_error(self): "The following user classes have the same class name: locust.test.fake_module1_for_env_test.MyUserWithSameName, locust.test.fake_module2_for_env_test.MyUserWithSameName", ) - def test_assign_same_task_weights(self): + def test_assign_uniform_weights(self): def verify_tasks(u, target_tasks): self.assertEqual(len(u.tasks), len(target_tasks)) tasks = [t.__name__ for t in u.tasks] @@ -61,7 +61,7 @@ def my_task_2(self): pass environment = Environment(user_classes=[MyUser1]) - environment.assign_same_task_weights() + environment.assign_uniform_weights() u = environment.user_classes[0] verify_tasks(u, ["my_task", "my_task_2"]) @@ -84,7 +84,7 @@ def stats(self): self.client.get("/stats/requests") environment = Environment(user_classes=[MyUser2]) - environment.assign_same_task_weights() + environment.assign_uniform_weights() u = environment.user_classes[0] verify_tasks(u, ["index", "stop", "stats"]) @@ -102,7 +102,7 @@ class MyUser3(User): tasks = [SingleTaskSet, outside_task] environment = Environment(user_classes=[MyUser3]) - environment.assign_same_task_weights() + environment.assign_uniform_weights() u = environment.user_classes[0] verify_tasks(u, ["outside_task", "outside_task_2"]) @@ -126,7 +126,20 @@ def dict_task_3(): class MyUser4(User): tasks = [DictTaskSet, SingleTaskSet, SingleTaskSet] + # Assign user tasks in dict environment = Environment(user_classes=[MyUser4]) - environment.assign_same_task_weights() + environment.assign_uniform_weights() + u = environment.user_classes[0] + verify_tasks(u, ["outside_task", "outside_task_2", "dict_task_1", "dict_task_2", "dict_task_3"]) + + class MyUser5(User): + tasks = { + DictTaskSet: 5, + SingleTaskSet: 3, + outside_task: 6, + } + + environment = Environment(user_classes=[MyUser5]) + environment.assign_uniform_weights() u = environment.user_classes[0] verify_tasks(u, ["outside_task", "outside_task_2", "dict_task_1", "dict_task_2", "dict_task_3"]) From f896f73d6fdf4cf77658732dbe3fa846fd9f5577 Mon Sep 17 00:00:00 2001 From: Shekar Ramaswamy Date: Mon, 9 Aug 2021 09:31:20 -0700 Subject: [PATCH 6/7] Changed flag to --equal-weights --- locust/argument_parser.py | 6 +++--- locust/main.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/locust/argument_parser.py b/locust/argument_parser.py index 1503811425..22817a898e 100644 --- a/locust/argument_parser.py +++ b/locust/argument_parser.py @@ -418,11 +418,11 @@ def setup_parser_arguments(parser): env_var="LOCUST_STOP_TIMEOUT", ) other_group.add_argument( - "--same-weights", + "--equal-weights", action="store_true", default=False, - dest="same_weights", - help="Use uniformly distributed task weights, overriding the weights specified in the locustfile.", + dest="equal_weights", + help="Use equally distributed task weights, overriding the weights specified in the locustfile.", ) user_classes_group = parser.add_argument_group("User classes") diff --git a/locust/main.py b/locust/main.py index e609a4a953..c41976b140 100644 --- a/locust/main.py +++ b/locust/main.py @@ -301,7 +301,7 @@ def main(): def assign_uniform_weights(environment, **kwargs): environment.assign_uniform_weights() - if options.same_weights: + if options.equal_weights: environment.events.init.add_listener(assign_uniform_weights) # Fire locust init event which can be used by end-users' code to run setup code that From f06db762abfa0d98230609e41a79d3cb98ae065d Mon Sep 17 00:00:00 2001 From: Shekar Ramaswamy Date: Mon, 9 Aug 2021 11:28:32 -0700 Subject: [PATCH 7/7] Update naming --- locust/env.py | 2 +- locust/main.py | 6 +++--- locust/test/test_env.py | 12 ++++++------ 3 files changed, 10 insertions(+), 10 deletions(-) diff --git a/locust/env.py b/locust/env.py index 101264bcc0..51114de0f5 100644 --- a/locust/env.py +++ b/locust/env.py @@ -216,7 +216,7 @@ def _filter_tasks_by_tags(self): for user_class in self.user_classes: filter_tasks_by_tags(user_class, self.tags, self.exclude_tags) - def assign_uniform_weights(self): + def assign_equal_weights(self): """ Update the user classes such that each user runs their specified tasks with equal probability. diff --git a/locust/main.py b/locust/main.py index c41976b140..1288d1e556 100644 --- a/locust/main.py +++ b/locust/main.py @@ -298,11 +298,11 @@ def main(): else: web_ui = None - def assign_uniform_weights(environment, **kwargs): - environment.assign_uniform_weights() + def assign_equal_weights(environment, **kwargs): + environment.assign_equal_weights() if options.equal_weights: - environment.events.init.add_listener(assign_uniform_weights) + environment.events.init.add_listener(assign_equal_weights) # Fire locust init event which can be used by end-users' code to run setup code that # need access to the Environment, Runner or WebUI. diff --git a/locust/test/test_env.py b/locust/test/test_env.py index 84558e5960..dccaad1527 100644 --- a/locust/test/test_env.py +++ b/locust/test/test_env.py @@ -41,7 +41,7 @@ def test_user_classes_with_same_name_is_error(self): "The following user classes have the same class name: locust.test.fake_module1_for_env_test.MyUserWithSameName, locust.test.fake_module2_for_env_test.MyUserWithSameName", ) - def test_assign_uniform_weights(self): + def test_assign_equal_weights(self): def verify_tasks(u, target_tasks): self.assertEqual(len(u.tasks), len(target_tasks)) tasks = [t.__name__ for t in u.tasks] @@ -61,7 +61,7 @@ def my_task_2(self): pass environment = Environment(user_classes=[MyUser1]) - environment.assign_uniform_weights() + environment.assign_equal_weights() u = environment.user_classes[0] verify_tasks(u, ["my_task", "my_task_2"]) @@ -84,7 +84,7 @@ def stats(self): self.client.get("/stats/requests") environment = Environment(user_classes=[MyUser2]) - environment.assign_uniform_weights() + environment.assign_equal_weights() u = environment.user_classes[0] verify_tasks(u, ["index", "stop", "stats"]) @@ -102,7 +102,7 @@ class MyUser3(User): tasks = [SingleTaskSet, outside_task] environment = Environment(user_classes=[MyUser3]) - environment.assign_uniform_weights() + environment.assign_equal_weights() u = environment.user_classes[0] verify_tasks(u, ["outside_task", "outside_task_2"]) @@ -128,7 +128,7 @@ class MyUser4(User): # Assign user tasks in dict environment = Environment(user_classes=[MyUser4]) - environment.assign_uniform_weights() + environment.assign_equal_weights() u = environment.user_classes[0] verify_tasks(u, ["outside_task", "outside_task_2", "dict_task_1", "dict_task_2", "dict_task_3"]) @@ -140,6 +140,6 @@ class MyUser5(User): } environment = Environment(user_classes=[MyUser5]) - environment.assign_uniform_weights() + environment.assign_equal_weights() u = environment.user_classes[0] verify_tasks(u, ["outside_task", "outside_task_2", "dict_task_1", "dict_task_2", "dict_task_3"])