Merge pull request #2393 from noirbizarre/shape/inheritance

feat(shape): support user base custom load shape

docs/custom-load-shape.rst

@@ -41,6 +41,18 @@ This functionality is further demonstrated in the `examples on github <https://g
 
 One further method may be helpful for your custom load shapes: `get_current_user_count()`, which returns the total number of active users. This method can be used to prevent advancing to subsequent steps until the desired number of users has been reached. This is especially useful if the initialization process for each user is slow or erratic in how long it takes. If this sounds like your use case, see the `example on github <https://github.com/locustio/locust/tree/master/examples/custom_shape/wait_user_count.py>`_.
 
+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.
+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.

locust/shape.py

@@ -1,20 +1,33 @@
 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.
     """
 
     runner: Optional[Runner] = None
     """Reference to the :class:`Runner <locust.runners.Runner>` instance"""
 
+    abstract: ClassVar[bool] = True
+
     def __init__(self):
         self.start_time = time.perf_counter()
 

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=[

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]]: