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)