From 38773c805961f9eedecab68d5e535c75f55c7c4e Mon Sep 17 00:00:00 2001 From: Chenyang Yan Date: Wed, 26 Jan 2022 22:23:52 +0800 Subject: [PATCH] hardening Environment.shape_class for distinct usage Hardening for issue 1981, make it more clear if the user made a mistake. Signed-off-by: Chenyang Yan --- locust/env.py | 5 +++++ locust/test/test_env.py | 11 ++++++++++- 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/locust/env.py b/locust/env.py index b5e9f68ede..a51f9e326c 100644 --- a/locust/env.py +++ b/locust/env.py @@ -98,6 +98,11 @@ def __init__( ", ".join(map(methodcaller("fullname"), self.user_classes)) ) ) + if self.shape_class is not None and not isinstance(self.shape_class, LoadTestShape): + raise ValueError( + "shape_class should be instance of LoadTestShape or subclass LoadTestShape, but got: %s" + % self.shape_class + ) def _create_runner( self, diff --git a/locust/test/test_env.py b/locust/test/test_env.py index 09ba4e97a3..2db8d60171 100644 --- a/locust/test/test_env.py +++ b/locust/test/test_env.py @@ -1,7 +1,7 @@ from locust import ( constant, ) -from locust.env import Environment +from locust.env import Environment, LoadTestShape from locust.user import ( User, task, @@ -190,3 +190,12 @@ def my_task(self): e.exception.args[0], "There are no users with weight > 0.", ) + + def test_shape_class_attribute(self): + class SubLoadTestShape(LoadTestShape): + """Inherited from locust.env.LoadTestShape""" + + with self.assertRaisesRegex( + ValueError, r"instance of LoadTestShape or subclass LoadTestShape", msg="exception message is mismatching" + ): + Environment(user_classes=[MyUserWithSameName1], shape_class=SubLoadTestShape)