diff --git a/docs/custom-load-shape.rst b/docs/custom-load-shape.rst index 2bb013c0a2..d767b51965 100644 --- a/docs/custom-load-shape.rst +++ b/docs/custom-load-shape.rst @@ -41,6 +41,18 @@ This functionality is further demonstrated in the `examples on github `_. +Note that if you want to defined your own custom base shape, you need to define the `abstract` attribute to `True` to avoid it being picked as Shape when imported: + +.. code-block:: python + + class MyBaseShape(LoadTestShape): + abstract = True + + def tick(self): + # Something reusable but needing inheritance + return None + + Combining Users with different load profiles -------------------------------------------- @@ -80,4 +92,4 @@ Adding the element ``user_classes`` to the return value gives you more detailed return None -This shape would create create in the first 10 seconds 10 User of ``UserA``. In the next twenty seconds 40 of type ``UserA / UserB`` and this continues until the stages end. \ No newline at end of file +This shape would create create in the first 10 seconds 10 User of ``UserA``. In the next twenty seconds 40 of type ``UserA / UserB`` and this continues until the stages end. diff --git a/locust/shape.py b/locust/shape.py index 5e09036100..afdeccd8fd 100644 --- a/locust/shape.py +++ b/locust/shape.py @@ -1,13 +1,24 @@ from __future__ import annotations import time -from typing import Optional, Tuple, List, Type -from abc import ABC, abstractmethod +from typing import ClassVar, Optional, Tuple, List, Type +from abc import ABCMeta, abstractmethod from . import User from .runners import Runner -class LoadTestShape(ABC): +class LoadTestShapeMeta(ABCMeta): + """ + Meta class for the main User class. It's used to allow User classes to specify task execution + ratio using an {task:int} dict, or a [(task0,int), ..., (taskN,int)] list. + """ + + def __new__(mcs, classname, bases, class_dict): + class_dict["abstract"] = class_dict.get("abstract", False) + return super().__new__(mcs, classname, bases, class_dict) + + +class LoadTestShape(metaclass=LoadTestShapeMeta): """ Base class for custom load shapes. """ @@ -15,6 +26,8 @@ class LoadTestShape(ABC): runner: Optional[Runner] = None """Reference to the :class:`Runner ` instance""" + abstract: ClassVar[bool] = True + def __init__(self): self.start_time = time.perf_counter() diff --git a/locust/test/test_load_locustfile.py b/locust/test/test_load_locustfile.py index f1d2082f4c..79aa502c44 100644 --- a/locust/test/test_load_locustfile.py +++ b/locust/test/test_load_locustfile.py @@ -76,6 +76,41 @@ def tick(self): self.assertNotIn("NotUserSubclass", user_classes) self.assertEqual(shape_class.__class__.__name__, "LoadTestShape") + def test_with_abstract_shape_class(self): + content = MOCK_LOCUSTFILE_CONTENT + textwrap.dedent( + """\ + class UserBaseLoadTestShape(LoadTestShape): + abstract = True + + def tick(self): + pass + + + class UserLoadTestShape(UserBaseLoadTestShape): + pass + """ + ) + + with mock_locustfile(content=content) as mocked: + _, user_classes, shape_class = main.load_locustfile(mocked.file_path) + self.assertNotIn("UserBaseLoadTestShape", user_classes) + self.assertNotIn("UserLoadTestShape", user_classes) + self.assertEqual(shape_class.__class__.__name__, "UserLoadTestShape") + + def test_with_not_imported_shape_class(self): + content = MOCK_LOCUSTFILE_CONTENT + textwrap.dedent( + """\ + class UserLoadTestShape(LoadTestShape): + def tick(self): + pass + """ + ) + + with mock_locustfile(content=content) as mocked: + _, user_classes, shape_class = main.load_locustfile(mocked.file_path) + self.assertNotIn("UserLoadTestShape", user_classes) + self.assertEqual(shape_class.__class__.__name__, "UserLoadTestShape") + def test_create_environment(self): options = parse_options( args=[ diff --git a/locust/util/load_locustfile.py b/locust/util/load_locustfile.py index 878cecd541..6eae8d108b 100644 --- a/locust/util/load_locustfile.py +++ b/locust/util/load_locustfile.py @@ -18,9 +18,7 @@ def is_shape_class(item): """ Check if a class is a LoadTestShape """ - return bool( - inspect.isclass(item) and issubclass(item, LoadTestShape) and item.__dict__["__module__"] != "locust.shape" - ) + return bool(inspect.isclass(item) and issubclass(item, LoadTestShape) and not getattr(item, "abstract", True)) def load_locustfile(path) -> Tuple[Optional[str], Dict[str, User], Optional[LoadTestShape]]: