From 673db60885577a0c0a4b876d4ec64bde9e583d4e Mon Sep 17 00:00:00 2001
From: Max Williams <max.williams@deliveryhero.com>
Date: Wed, 29 Jul 2020 12:35:36 +0200
Subject: [PATCH 01/19] Initial commit with shaper

---
 examples/custom_shape_stages.py | 35 ++++++++++++++
 locust/__init__.py              |  1 +
 locust/env.py                   |  8 +++-
 locust/main.py                  | 34 +++++++++++---
 locust/runners.py               | 33 ++++++++++++++
 locust/shapers.py               | 81 +++++++++++++++++++++++++++++++++
 locust/web.py                   |  6 ++-
 7 files changed, 190 insertions(+), 8 deletions(-)
 create mode 100644 examples/custom_shape_stages.py
 create mode 100644 locust/shapers.py

diff --git a/examples/custom_shape_stages.py b/examples/custom_shape_stages.py
new file mode 100644
index 0000000000..04ca323fad
--- /dev/null
+++ b/examples/custom_shape_stages.py
@@ -0,0 +1,35 @@
+from locust import HttpUser, TaskSet, task, constant
+from locust import StagesShaper
+
+
+class UserTasks(TaskSet):
+    @task
+    def one(self):
+        self.client.get("/one")
+
+    @task
+    def two(self):
+        self.client.get("/two")
+
+    @task
+    def three(self):
+        self.client.get("/three")
+
+
+class WebsiteUser(HttpUser):
+    wait_time = constant(0.5)
+    tasks = [UserTasks]
+
+
+class myShape(StagesShaper):
+    def __init__(self):
+        self.stages = [
+            {'duration': 60, 'users': 10, 'hatch_rate': 10, 'stop': False},
+            {'duration': 80, 'users': 30, 'hatch_rate': 10, 'stop': False},
+            {'duration': 100, 'users': 50, 'hatch_rate': 10, 'stop': False},
+            {'duration': 120, 'users': 70, 'hatch_rate': 10, 'stop': False},
+            {'duration': 180, 'users': 100, 'hatch_rate': 10, 'stop': False},
+            {'duration': 300, 'users': 10, 'hatch_rate': 10, 'stop': False},
+            {'duration': 320, 'users': 5, 'hatch_rate': 10, 'stop': False},
+            {'duration': 360, 'users': 1, 'hatch_rate': 1, 'stop': True},
+        ]
diff --git a/locust/__init__.py b/locust/__init__.py
index 4885f97632..75747064ff 100644
--- a/locust/__init__.py
+++ b/locust/__init__.py
@@ -8,6 +8,7 @@
 from .user.task import task, tag, TaskSet
 from .user.users import HttpUser, User
 from .user.wait_time import between, constant, constant_pacing
+from .shapers import StagesShaper, StepLoadShaper
 
 from .event import Events
 events = Events()
diff --git a/locust/env.py b/locust/env.py
index 1576245575..6f61c03316 100644
--- a/locust/env.py
+++ b/locust/env.py
@@ -4,6 +4,7 @@
 from .runners import LocalRunner, MasterRunner, WorkerRunner
 from .web import WebUI
 from .user.task import filter_tasks_by_tags
+from .shapers import LoadTestShaper
 
 
 class Environment:
@@ -15,7 +16,10 @@ class Environment:
     
     user_classes = []
     """User classes that the runner will run"""
-    
+
+    shaper_class = None
+    """shaper classes that the runner will run"""
+
     tags = None
     """If set, only tasks that are tagged by tags in this list will be executed"""
 
@@ -63,6 +67,7 @@ class Environment:
     def  __init__(
         self, *,
         user_classes=[],
+        shaper_class=None,
         tags=None,
         exclude_tags=None,
         events=None, 
@@ -79,6 +84,7 @@ def  __init__(
             self.events = Events()
         
         self.user_classes = user_classes
+        self.shaper_class = shaper_class
         self.tags = tags
         self.exclude_tags = exclude_tags
         self.stats = RequestStats()
diff --git a/locust/main.py b/locust/main.py
index 61a880924e..b39593ab0e 100644
--- a/locust/main.py
+++ b/locust/main.py
@@ -22,6 +22,9 @@
 from .util.timespan import parse_timespan
 from .exception import AuthCredentialsError
 
+from .shapers import LoadTestShaper
+
+
 version = locust.__version__
 
 
@@ -36,6 +39,16 @@ def is_user_class(item):
     )
 
 
+def is_shaper_class(item):
+    """
+    Check if a class is a LoadTestShaper
+    """
+    return bool(
+        inspect.isclass(item)
+        and issubclass(item, LoadTestShaper)
+        and item.__dict__['__module__'] != 'locust.shapers'
+    )
+
 def load_locustfile(path):
     """
     Import given locustfile path and return (docstring, callables).
@@ -84,15 +97,24 @@ def __import_locustfile__(filename, path):
         del sys.path[0]
     # Return our two-tuple
     user_classes = {name:value for name, value in vars(imported).items() if is_user_class(value)}
-    return imported.__doc__, user_classes
 
+    # Find shaper class, if any, return it
+    shaper_classes = [value for name, value in vars(imported).items() if is_shaper_class(value)]
+    if shaper_classes:
+        shaper_class = shaper_classes[0]()
+    else:
+        shaper_class = None
+
+    return imported.__doc__, user_classes, shaper_class
 
-def create_environment(user_classes, options, events=None):
+
+def create_environment(user_classes, shaper_class, options, events=None):
     """
     Create an Environment instance from options
     """
     return Environment(
         user_classes=user_classes,
+        shaper_class=shaper_class,
         tags=options.tags,
         exclude_tags=options.exclude_tags,
         events=events,
@@ -110,8 +132,8 @@ def main():
     locustfile = parse_locustfile_option()
     
     # import the locustfile
-    docstring, user_classes = load_locustfile(locustfile)
-    
+    docstring, user_classes, shaper_class = load_locustfile(locustfile)
+
     # parse all command line options
     options = parse_options()
 
@@ -163,8 +185,8 @@ def main():
         logger.warning("System open file limit setting is not high enough for load testing, and the OS didn't allow locust to increase it by itself. See https://docs.locust.io/en/stable/installation.html#increasing-maximum-number-of-open-files-limit for more info.")
 
     # create locust Environment
-    environment = create_environment(user_classes, options, events=locust.events)
-    
+    environment = create_environment(user_classes, shaper_class, options, events=locust.events)
+
     if options.show_task_ratio:
         print("\n Task ratio per User class")
         print( "-" * 80)
diff --git a/locust/runners.py b/locust/runners.py
index f9971bd126..2dde52768c 100644
--- a/locust/runners.py
+++ b/locust/runners.py
@@ -49,6 +49,8 @@ def __init__(self, environment):
         self.state = STATE_INIT
         self.hatching_greenlet = None
         self.stepload_greenlet = None
+        self.shaper_greenlet = None
+        self.shaper_last_state = None
         self.current_cpu_usage = 0
         self.cpu_warning_emitted = False
         self.greenlet.spawn(self.monitor_cpu).link_exception(greenlet_exception_handler)
@@ -245,6 +247,8 @@ def start(self, user_count, hatch_rate, wait=False):
                      If False (the default), a greenlet that spawns the users will be 
                      started and the call to this method will return immediately.
         """
+        print('Runner start')
+
         if self.state != STATE_RUNNING and self.state != STATE_HATCHING:
             self.stats.clear_all()
             self.exceptions = {}
@@ -294,6 +298,35 @@ def stepload_worker(self, hatch_rate, step_users_growth, step_duration):
             logger.info('Step loading: start hatch job of %d user.' % (current_num_users))
             gevent.sleep(step_duration)
 
+    def start_shaper(self):
+        print('start_shaper')
+        if self.shaper_greenlet:
+            logger.info("There is an ongoing shaper running, will stop it now.")
+            self.shaper_greenlet.kill()
+
+        logger.info("Starting a new shaper load test")
+        self.state = STATE_INIT
+        self.shaper_greenlet = self.greenlet.spawn(self.shaper_worker)
+        self.shaper_greenlet.link_exception(greenlet_exception_handler)
+
+
+    def shaper_worker(self):
+        print('Shaper worker starting')
+        while self.state == STATE_INIT or self.state == STATE_HATCHING or self.state == STATE_RUNNING:
+            new_state = self.environment.shaper_class.tick()
+            user_count, hatch_rate, stop_test = new_state
+            if stop_test:
+                print('Shaper test finished')
+                self.stop()
+            if self.shaper_last_state == new_state:
+                gevent.sleep(1)
+            else:
+                new_user_count = user_count - self.user_count
+                print('Adding {0} users'.format(new_user_count))
+                self.start(user_count=new_user_count, hatch_rate=hatch_rate)
+                self.shaper_last_state = new_state
+                print('Custom setting applied: {0}'.format(self.shaper_last_state))
+
     def stop(self):
         """
         Stop a running load test by stopping all running users
diff --git a/locust/shapers.py b/locust/shapers.py
new file mode 100644
index 0000000000..c7f1b13551
--- /dev/null
+++ b/locust/shapers.py
@@ -0,0 +1,81 @@
+import time
+import math
+
+
+class LoadTestShaper(object):
+    '''
+    A simple load test shaper class used to control the shape of load generated
+    during a load test.
+    '''
+
+    start_time = time.monotonic()
+
+    def tick(self):
+        '''
+        Returns a tuple with 3 elements to control the running load test.
+
+            user_count: Total user count
+            hatch_rate: Hatch rate to use when changing total user count
+            stop_test: A boolean to stop the test
+
+        '''
+        return (0, 0, True)
+
+
+class StagesShaper(LoadTestShaper):
+    '''
+    A simply load test shaper class that can be passed a list of dicts that
+    represent stages. Each stage has these keys:
+
+        duration - When this many seconds pass the test is advanced to the next stage
+        users - Total user count
+        hatch_rate - Hatch rate
+        stop - A boolean that can stop that test at a specific stage
+
+    Arguments:
+
+        stop_at_end - can be set to stop once all stages have run.
+    '''
+    def __init__(self, stages, stop_at_end=False):
+        self.stages = sorted(stages, key=lambda k: k['duration'])
+        self.stop_at_end = stop_at_end
+
+    def tick(self):
+        run_time = round(time.monotonic() - self.start_time)
+
+        for stage in self.stages:
+            if run_time < stage['duration']:
+                tick_data = (stage['users'], stage['hatch_rate'], stage['stop'])
+                self.last_stage = tick_data
+                return tick_data
+
+        if self.stop_at_end:
+            print('Stopping test')
+            return (0, 0, True)
+        else:
+            print('Continuing test')
+            return self.last_stage
+
+
+
+class StepLoadShaper(LoadTestShaper):
+    '''
+    A step load shaper.
+
+    Arguments:
+
+        step_time - Time between steps
+        step_load - User increase amount at each step
+        hatch_rate - Hatch rate to use at every step
+        time_limit - Time limit in seconds
+    '''
+    def __init__(self, step_time, step_load, time_limit):
+        self.step_time = step_time
+        self.step_load = step_load
+        self.hatch_rate = hatch_rate
+        self.time_limit = time_limit
+
+    def tick(self):
+        run_time = round(time.monotonic() - self.start_time)
+        current_step = math.floor(run_time / step_time) + 1
+        return (current_step * self.step_load, self.hatch_rate, run_time > self.time_limit)
diff --git a/locust/web.py b/locust/web.py
index 145f3d9ce6..61f3648fae 100644
--- a/locust/web.py
+++ b/locust/web.py
@@ -152,7 +152,11 @@ def swarm():
                 step_duration = parse_timespan(str(request.form["step_duration"]))
                 environment.runner.start_stepload(user_count, hatch_rate, step_user_count, step_duration)
                 return jsonify({'success': True, 'message': 'Swarming started in Step Load Mode', 'host': environment.host})
-            
+
+            if environment.shaper_class:
+                environment.runner.start_shaper()
+                return jsonify({'success': True, 'message': 'Swarming started with custom shaper', 'host': environment.host})
+
             environment.runner.start(user_count, hatch_rate)
             return jsonify({'success': True, 'message': 'Swarming started', 'host': environment.host})
         

From fa4c70dbedd0cc5cadb11c1403bb7b384b15abc3 Mon Sep 17 00:00:00 2001
From: Max Williams <max.williams@deliveryhero.com>
Date: Fri, 31 Jul 2020 16:38:01 +0200
Subject: [PATCH 02/19] renaming to shape, add some logging, examples

---
 examples/custom_shape/double_wave.py | 55 +++++++++++++++++++
 examples/custom_shape/stages.py      | 58 ++++++++++++++++++++
 examples/custom_shape/step_load.py   | 42 +++++++++++++++
 examples/custom_shape_stages.py      | 35 ------------
 locust/__init__.py                   |  2 +-
 locust/env.py                        | 18 +++----
 locust/main.py                       | 31 ++++++-----
 locust/runners.py                    | 46 ++++++++--------
 locust/shape.py                      | 30 +++++++++++
 locust/shapers.py                    | 81 ----------------------------
 locust/web.py                        |  6 +--
 11 files changed, 235 insertions(+), 169 deletions(-)
 create mode 100644 examples/custom_shape/double_wave.py
 create mode 100644 examples/custom_shape/stages.py
 create mode 100644 examples/custom_shape/step_load.py
 delete mode 100644 examples/custom_shape_stages.py
 create mode 100644 locust/shape.py
 delete mode 100644 locust/shapers.py

diff --git a/examples/custom_shape/double_wave.py b/examples/custom_shape/double_wave.py
new file mode 100644
index 0000000000..72ede36561
--- /dev/null
+++ b/examples/custom_shape/double_wave.py
@@ -0,0 +1,55 @@
+import math
+from locust import HttpUser, TaskSet, task, constant
+from locust import LoadTestShape
+
+
+class UserTasks(TaskSet):
+    @task
+    def get_root(self):
+        self.client.get("/")
+
+
+class WebsiteUser(HttpUser):
+    wait_time = constant(0.5)
+    tasks = [UserTasks]
+
+
+class DoubleWave(LoadTestShape):
+    '''
+    A shape to immitate some specific user behaviour. In this example, midday
+    and evening meal times.
+
+    Arguments:
+
+        min_users - minimum users
+        peak_one - first peak size
+        peak_two - second peak size
+        amplitude - size of the test
+        time_limit - total length of test
+    '''
+    def __init__(self,
+            min_users=20,
+            peak_one=1,
+            peak_two=1.5,
+            amplitude=150,
+            time_limit=600,
+        ):
+        self.min_users = min_users
+        self.peak_one = peak_one
+        self.peak_two = peak_two
+        self.amplitude = amplitude
+        self.time_limit = time_limit
+
+        self.first_peak_time = round(self.time_limit / 3)
+        self.first_peak_users = round(self.peak_one * self.amplitude)
+        self.second_peak_time = round(self.first_peak_time * 2)
+        self.second_peak_users = round(self.peak_two * self.amplitude)
+
+    def tick(self):
+        run_time = self.get_run_time()
+
+        if run_time < self.time_limit:
+            user_count = self.second_peak_users * math.e** - ((run_time/(math.sqrt(self.time_limit)*2))-4)**2 + self.first_peak_users * math.e** - ((run_time/(math.sqrt(self.time_limit)*2))-8)**2 + self.min_users
+            return (round(user_count), round(user_count), False)
+        else:
+            return (0, 0, True)
diff --git a/examples/custom_shape/stages.py b/examples/custom_shape/stages.py
new file mode 100644
index 0000000000..f4fc4216e8
--- /dev/null
+++ b/examples/custom_shape/stages.py
@@ -0,0 +1,58 @@
+from locust import HttpUser, TaskSet, task, constant
+from locust import LoadTestShape
+
+
+class UserTasks(TaskSet):
+    @task
+    def get_root(self):
+        self.client.get("/")
+
+
+class WebsiteUser(HttpUser):
+    wait_time = constant(0.5)
+    tasks = [UserTasks]
+
+
+class StagesShape(LoadTestShape):
+    '''
+    A simply load test shape class that has different user and hatch_rate at
+    different stages.
+
+        duration - When this many seconds pass the test is advanced to the next stage
+        users - Total user count
+        hatch_rate - Hatch rate
+        stop - A boolean that can stop that test at a specific stage
+
+    Arguments:
+
+        stop_at_end - can be set to stop once all stages have run.
+    '''
+    def __init__(self,
+            stages = [
+                {'duration': 60, 'users': 10, 'hatch_rate': 10},
+                {'duration': 80, 'users': 30, 'hatch_rate': 10},
+                {'duration': 100, 'users': 50, 'hatch_rate': 10},
+                {'duration': 120, 'users': 70, 'hatch_rate': 10},
+                {'duration': 180, 'users': 100, 'hatch_rate': 10},
+                {'duration': 220, 'users': 10, 'hatch_rate': 10},
+                {'duration': 230, 'users': 5, 'hatch_rate': 10},
+                {'duration': 240, 'users': 1, 'hatch_rate': 1},
+            ],
+            stop_at_end=True
+        ):
+        self.stages = sorted(stages, key=lambda k: k['duration'])
+        self.stop_at_end = stop_at_end
+
+    def tick(self):
+        run_time = self.get_run_time()
+
+        for stage in self.stages:
+            if run_time < stage['duration']:
+                tick_data = (stage['users'], stage['hatch_rate'], stage.get('stop', False))
+                self.last_stage = tick_data
+                return tick_data
+
+        if self.stop_at_end:
+            return (0, 0, True)
+        else:
+            return self.last_stage
diff --git a/examples/custom_shape/step_load.py b/examples/custom_shape/step_load.py
new file mode 100644
index 0000000000..51ca14a8d4
--- /dev/null
+++ b/examples/custom_shape/step_load.py
@@ -0,0 +1,42 @@
+import math
+from locust import HttpUser, TaskSet, task, constant
+from locust import LoadTestShape
+
+
+class UserTasks(TaskSet):
+    @task
+    def get_root(self):
+        self.client.get("/")
+
+
+class WebsiteUser(HttpUser):
+    wait_time = constant(0.5)
+    tasks = [UserTasks]
+
+
+class StepLoadShape(LoadTestShape):
+    '''
+    A step load shape
+
+    Arguments:
+
+        step_time - Time between steps
+        step_load - User increase amount at each step
+        hatch_rate - Hatch rate to use at every step
+        time_limit - Time limit in seconds
+    '''
+    def __init__(self,
+            step_time=30,
+            step_load=10,
+            hatch_rate=10,
+            time_limit=600
+        ):
+        self.step_time = step_time
+        self.step_load = step_load
+        self.hatch_rate = hatch_rate
+        self.time_limit = time_limit
+
+    def tick(self):
+        run_time = self.get_run_time()
+        current_step = math.floor(run_time / self.step_time) + 1
+        return (current_step * self.step_load, self.hatch_rate, run_time > self.time_limit)
diff --git a/examples/custom_shape_stages.py b/examples/custom_shape_stages.py
deleted file mode 100644
index 04ca323fad..0000000000
--- a/examples/custom_shape_stages.py
+++ /dev/null
@@ -1,35 +0,0 @@
-from locust import HttpUser, TaskSet, task, constant
-from locust import StagesShaper
-
-
-class UserTasks(TaskSet):
-    @task
-    def one(self):
-        self.client.get("/one")
-
-    @task
-    def two(self):
-        self.client.get("/two")
-
-    @task
-    def three(self):
-        self.client.get("/three")
-
-
-class WebsiteUser(HttpUser):
-    wait_time = constant(0.5)
-    tasks = [UserTasks]
-
-
-class myShape(StagesShaper):
-    def __init__(self):
-        self.stages = [
-            {'duration': 60, 'users': 10, 'hatch_rate': 10, 'stop': False},
-            {'duration': 80, 'users': 30, 'hatch_rate': 10, 'stop': False},
-            {'duration': 100, 'users': 50, 'hatch_rate': 10, 'stop': False},
-            {'duration': 120, 'users': 70, 'hatch_rate': 10, 'stop': False},
-            {'duration': 180, 'users': 100, 'hatch_rate': 10, 'stop': False},
-            {'duration': 300, 'users': 10, 'hatch_rate': 10, 'stop': False},
-            {'duration': 320, 'users': 5, 'hatch_rate': 10, 'stop': False},
-            {'duration': 360, 'users': 1, 'hatch_rate': 1, 'stop': True},
-        ]
diff --git a/locust/__init__.py b/locust/__init__.py
index 75747064ff..d6ec26e30d 100644
--- a/locust/__init__.py
+++ b/locust/__init__.py
@@ -8,7 +8,7 @@
 from .user.task import task, tag, TaskSet
 from .user.users import HttpUser, User
 from .user.wait_time import between, constant, constant_pacing
-from .shapers import StagesShaper, StepLoadShaper
+from .shape import LoadTestShape
 
 from .event import Events
 events = Events()
diff --git a/locust/env.py b/locust/env.py
index 6f61c03316..76e939f90f 100644
--- a/locust/env.py
+++ b/locust/env.py
@@ -4,7 +4,7 @@
 from .runners import LocalRunner, MasterRunner, WorkerRunner
 from .web import WebUI
 from .user.task import filter_tasks_by_tags
-from .shapers import LoadTestShaper
+from .shape import LoadTestShape
 
 
 class Environment:
@@ -17,8 +17,8 @@ class Environment:
     user_classes = []
     """User classes that the runner will run"""
 
-    shaper_class = None
-    """shaper classes that the runner will run"""
+    shape_class = None
+    """A shape class to control the shape of the load test"""
 
     tags = None
     """If set, only tasks that are tagged by tags in this list will be executed"""
@@ -67,13 +67,13 @@ class Environment:
     def  __init__(
         self, *,
         user_classes=[],
-        shaper_class=None,
+        shape_class=None,
         tags=None,
         exclude_tags=None,
-        events=None, 
-        host=None, 
-        reset_stats=False, 
-        step_load=False, 
+        events=None,
+        host=None,
+        reset_stats=False,
+        step_load=False,
         stop_timeout=None,
         catch_exceptions=True,
         parsed_options=None,
@@ -84,7 +84,7 @@ def  __init__(
             self.events = Events()
         
         self.user_classes = user_classes
-        self.shaper_class = shaper_class
+        self.shape_class = shape_class
         self.tags = tags
         self.exclude_tags = exclude_tags
         self.stats = RequestStats()
diff --git a/locust/main.py b/locust/main.py
index b39593ab0e..16b23db603 100644
--- a/locust/main.py
+++ b/locust/main.py
@@ -21,8 +21,7 @@
 from .user.inspectuser import get_task_ratio_dict, print_task_ratio
 from .util.timespan import parse_timespan
 from .exception import AuthCredentialsError
-
-from .shapers import LoadTestShaper
+from .shape import LoadTestShape
 
 
 version = locust.__version__
@@ -39,14 +38,14 @@ def is_user_class(item):
     )
 
 
-def is_shaper_class(item):
+def is_shape_class(item):
     """
-    Check if a class is a LoadTestShaper
+    Check if a class is a LoadTestShape
     """
     return bool(
         inspect.isclass(item)
-        and issubclass(item, LoadTestShaper)
-        and item.__dict__['__module__'] != 'locust.shapers'
+        and issubclass(item, LoadTestShape)
+        and item.__dict__['__module__'] != 'locust.shape'
     )
 
 def load_locustfile(path):
@@ -98,23 +97,23 @@ def __import_locustfile__(filename, path):
     # Return our two-tuple
     user_classes = {name:value for name, value in vars(imported).items() if is_user_class(value)}
 
-    # Find shaper class, if any, return it
-    shaper_classes = [value for name, value in vars(imported).items() if is_shaper_class(value)]
-    if shaper_classes:
-        shaper_class = shaper_classes[0]()
+    # Find shape class, if any, return it
+    shape_classes = [value for name, value in vars(imported).items() if is_shape_class(value)]
+    if shape_classes:
+        shape_class = shape_classes[0]()
     else:
-        shaper_class = None
+        shape_class = None
 
-    return imported.__doc__, user_classes, shaper_class
+    return imported.__doc__, user_classes, shape_class
 
 
-def create_environment(user_classes, shaper_class, options, events=None):
+def create_environment(user_classes, shape_class, options, events=None):
     """
     Create an Environment instance from options
     """
     return Environment(
         user_classes=user_classes,
-        shaper_class=shaper_class,
+        shape_class=shape_class,
         tags=options.tags,
         exclude_tags=options.exclude_tags,
         events=events,
@@ -132,7 +131,7 @@ def main():
     locustfile = parse_locustfile_option()
     
     # import the locustfile
-    docstring, user_classes, shaper_class = load_locustfile(locustfile)
+    docstring, user_classes, shape_class = load_locustfile(locustfile)
 
     # parse all command line options
     options = parse_options()
@@ -185,7 +184,7 @@ def main():
         logger.warning("System open file limit setting is not high enough for load testing, and the OS didn't allow locust to increase it by itself. See https://docs.locust.io/en/stable/installation.html#increasing-maximum-number-of-open-files-limit for more info.")
 
     # create locust Environment
-    environment = create_environment(user_classes, shaper_class, options, events=locust.events)
+    environment = create_environment(user_classes, shape_class, options, events=locust.events)
 
     if options.show_task_ratio:
         print("\n Task ratio per User class")
diff --git a/locust/runners.py b/locust/runners.py
index 2dde52768c..8933a3b1e6 100644
--- a/locust/runners.py
+++ b/locust/runners.py
@@ -49,8 +49,8 @@ def __init__(self, environment):
         self.state = STATE_INIT
         self.hatching_greenlet = None
         self.stepload_greenlet = None
-        self.shaper_greenlet = None
-        self.shaper_last_state = None
+        self.shape_greenlet = None
+        self.shape_last_state = None
         self.current_cpu_usage = 0
         self.cpu_warning_emitted = False
         self.greenlet.spawn(self.monitor_cpu).link_exception(greenlet_exception_handler)
@@ -73,10 +73,13 @@ def on_request_failure(request_type, name, response_time, response_length, excep
         def on_hatch_complete(user_count):
             self.state = STATE_RUNNING
             if environment.reset_stats:
-                logger.info("Resetting stats\n")
+                logger.info("Resetting stats")
                 self.stats.reset_all()
         self.environment.events.hatch_complete.add_listener(on_hatch_complete)
-    
+
+        if self.environment.shape_class:
+            logger.debug("Starting with shape class")
+
     def __del__(self):
         # don't leave any stray greenlets if runner is removed
         if self.greenlet and len(self.greenlet) > 0:
@@ -247,8 +250,6 @@ def start(self, user_count, hatch_rate, wait=False):
                      If False (the default), a greenlet that spawns the users will be 
                      started and the call to this method will return immediately.
         """
-        print('Runner start')
-
         if self.state != STATE_RUNNING and self.state != STATE_HATCHING:
             self.stats.clear_all()
             self.exceptions = {}
@@ -298,34 +299,31 @@ def stepload_worker(self, hatch_rate, step_users_growth, step_duration):
             logger.info('Step loading: start hatch job of %d user.' % (current_num_users))
             gevent.sleep(step_duration)
 
-    def start_shaper(self):
-        print('start_shaper')
-        if self.shaper_greenlet:
-            logger.info("There is an ongoing shaper running, will stop it now.")
-            self.shaper_greenlet.kill()
+    def start_shape(self):
+        logger.info("Shape test starting")
+        if self.shape_greenlet:
+            logger.info("There is an ongoing shape test running, will stop it now")
+            self.shape_greenlet.kill()
 
-        logger.info("Starting a new shaper load test")
         self.state = STATE_INIT
-        self.shaper_greenlet = self.greenlet.spawn(self.shaper_worker)
-        self.shaper_greenlet.link_exception(greenlet_exception_handler)
+        self.shape_greenlet = self.greenlet.spawn(self.shape_worker)
+        self.shape_greenlet.link_exception(greenlet_exception_handler)
 
 
-    def shaper_worker(self):
-        print('Shaper worker starting')
+    def shape_worker(self):
+        logger.info('Shape worker starting')
         while self.state == STATE_INIT or self.state == STATE_HATCHING or self.state == STATE_RUNNING:
-            new_state = self.environment.shaper_class.tick()
+            new_state = self.environment.shape_class.tick()
             user_count, hatch_rate, stop_test = new_state
             if stop_test:
-                print('Shaper test finished')
+                logger.info('Shape test stopping')
                 self.stop()
-            if self.shaper_last_state == new_state:
+            elif self.shape_last_state == new_state:
                 gevent.sleep(1)
             else:
-                new_user_count = user_count - self.user_count
-                print('Adding {0} users'.format(new_user_count))
-                self.start(user_count=new_user_count, hatch_rate=hatch_rate)
-                self.shaper_last_state = new_state
-                print('Custom setting applied: {0}'.format(self.shaper_last_state))
+                logger.info("Shape test updating to %d users at %.2f hatch rate" % (user_count, hatch_rate))
+                self.start(user_count=user_count, hatch_rate=hatch_rate)
+                self.shape_last_state = new_state
 
     def stop(self):
         """
diff --git a/locust/shape.py b/locust/shape.py
new file mode 100644
index 0000000000..efcea15ca0
--- /dev/null
+++ b/locust/shape.py
@@ -0,0 +1,30 @@
+import time
+
+
+class LoadTestShape(object):
+    '''
+    A simple load test shape class used to control the shape of load generated
+    during a load test.
+    '''
+
+    start_time = time.monotonic()
+
+    def get_run_time(self):
+        '''
+        Calculates run time in seconds of the load test
+        '''
+        return round(time.monotonic() - self.start_time)
+        
+
+    def tick(self):
+        '''
+        Returns a tuple with 3 elements to control the running load test.
+
+            user_count: Total user count
+            hatch_rate: Hatch rate to use when changing total user count
+            stop_test: A boolean to stop the test
+
+        '''
+        run_time = self.get_run_time()
+
+        return (0, 0, True)
diff --git a/locust/shapers.py b/locust/shapers.py
deleted file mode 100644
index c7f1b13551..0000000000
--- a/locust/shapers.py
+++ /dev/null
@@ -1,81 +0,0 @@
-import time
-import math
-
-
-class LoadTestShaper(object):
-    '''
-    A simple load test shaper class used to control the shape of load generated
-    during a load test.
-    '''
-
-    start_time = time.monotonic()
-
-    def tick(self):
-        '''
-        Returns a tuple with 3 elements to control the running load test.
-
-            user_count: Total user count
-            hatch_rate: Hatch rate to use when changing total user count
-            stop_test: A boolean to stop the test
-
-        '''
-        return (0, 0, True)
-
-
-class StagesShaper(LoadTestShaper):
-    '''
-    A simply load test shaper class that can be passed a list of dicts that
-    represent stages. Each stage has these keys:
-
-        duration - When this many seconds pass the test is advanced to the next stage
-        users - Total user count
-        hatch_rate - Hatch rate
-        stop - A boolean that can stop that test at a specific stage
-
-    Arguments:
-
-        stop_at_end - can be set to stop once all stages have run.
-    '''
-    def __init__(self, stages, stop_at_end=False):
-        self.stages = sorted(stages, key=lambda k: k['duration'])
-        self.stop_at_end = stop_at_end
-
-    def tick(self):
-        run_time = round(time.monotonic() - self.start_time)
-
-        for stage in self.stages:
-            if run_time < stage['duration']:
-                tick_data = (stage['users'], stage['hatch_rate'], stage['stop'])
-                self.last_stage = tick_data
-                return tick_data
-
-        if self.stop_at_end:
-            print('Stopping test')
-            return (0, 0, True)
-        else:
-            print('Continuing test')
-            return self.last_stage
-
-
-
-class StepLoadShaper(LoadTestShaper):
-    '''
-    A step load shaper.
-
-    Arguments:
-
-        step_time - Time between steps
-        step_load - User increase amount at each step
-        hatch_rate - Hatch rate to use at every step
-        time_limit - Time limit in seconds
-    '''
-    def __init__(self, step_time, step_load, time_limit):
-        self.step_time = step_time
-        self.step_load = step_load
-        self.hatch_rate = hatch_rate
-        self.time_limit = time_limit
-
-    def tick(self):
-        run_time = round(time.monotonic() - self.start_time)
-        current_step = math.floor(run_time / step_time) + 1
-        return (current_step * self.step_load, self.hatch_rate, run_time > self.time_limit)
diff --git a/locust/web.py b/locust/web.py
index 61f3648fae..1d73e9a41c 100644
--- a/locust/web.py
+++ b/locust/web.py
@@ -153,9 +153,9 @@ def swarm():
                 environment.runner.start_stepload(user_count, hatch_rate, step_user_count, step_duration)
                 return jsonify({'success': True, 'message': 'Swarming started in Step Load Mode', 'host': environment.host})
 
-            if environment.shaper_class:
-                environment.runner.start_shaper()
-                return jsonify({'success': True, 'message': 'Swarming started with custom shaper', 'host': environment.host})
+            if environment.shape_class:
+                environment.runner.start_shape()
+                return jsonify({'success': True, 'message': 'Swarming started using shape class', 'host': environment.host})
 
             environment.runner.start(user_count, hatch_rate)
             return jsonify({'success': True, 'message': 'Swarming started', 'host': environment.host})

From 1782959e636ed5c0b7bfe69625dcac7b42ae7a87 Mon Sep 17 00:00:00 2001
From: Max Williams <max.williams@deliveryhero.com>
Date: Wed, 5 Aug 2020 12:13:20 +0200
Subject: [PATCH 03/19] fix existing tests by adjusting param order

---
 locust/main.py           | 4 ++--
 locust/test/test_main.py | 8 ++++----
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/locust/main.py b/locust/main.py
index 16b23db603..04526b5d70 100644
--- a/locust/main.py
+++ b/locust/main.py
@@ -107,7 +107,7 @@ def __import_locustfile__(filename, path):
     return imported.__doc__, user_classes, shape_class
 
 
-def create_environment(user_classes, shape_class, options, events=None):
+def create_environment(user_classes, options, events=None, shape_class=None):
     """
     Create an Environment instance from options
     """
@@ -184,7 +184,7 @@ def main():
         logger.warning("System open file limit setting is not high enough for load testing, and the OS didn't allow locust to increase it by itself. See https://docs.locust.io/en/stable/installation.html#increasing-maximum-number-of-open-files-limit for more info.")
 
     # create locust Environment
-    environment = create_environment(user_classes, shape_class, options, events=locust.events)
+    environment = create_environment(user_classes, options, events=locust.events, shape_class=shape_class)
 
     if options.show_task_ratio:
         print("\n Task ratio per User class")
diff --git a/locust/test/test_main.py b/locust/test/test_main.py
index 78e3556117..3a8aed84d9 100644
--- a/locust/test/test_main.py
+++ b/locust/test/test_main.py
@@ -44,21 +44,21 @@ class ThriftLocust(User):
     
     def test_load_locust_file_from_absolute_path(self):
         with mock_locustfile() as mocked:
-            docstring, user_classes = main.load_locustfile(mocked.file_path)
+            docstring, user_classes, shape_class = main.load_locustfile(mocked.file_path)
             self.assertIn('UserSubclass', user_classes)
             self.assertNotIn('NotUserSubclass', user_classes)
 
     def test_load_locust_file_from_relative_path(self):
         with mock_locustfile() as mocked:
-            docstring, user_classes = main.load_locustfile(os.path.join('./locust/test/', mocked.filename))
+            docstring, user_classes, shape_class = main.load_locustfile(os.path.join('./locust/test/', mocked.filename))
 
     def test_load_locust_file_with_a_dot_in_filename(self):
         with mock_locustfile(filename_prefix="mocked.locust.file") as mocked:
-            docstring, user_classes = main.load_locustfile(mocked.file_path)
+            docstring, user_classes, shape_class = main.load_locustfile(mocked.file_path)
     
     def test_return_docstring_and_user_classes(self):
         with mock_locustfile() as mocked:
-            docstring, user_classes = main.load_locustfile(mocked.file_path)
+            docstring, user_classes, shape_class = main.load_locustfile(mocked.file_path)
             self.assertEqual("This is a mock locust file for unit testing", docstring)
             self.assertIn('UserSubclass', user_classes)
             self.assertNotIn('NotUserSubclass', user_classes)

From b5e9fdc0505a1a87c5f199608217dd6bc7ae4e5f Mon Sep 17 00:00:00 2001
From: Max Williams <max.williams@deliveryhero.com>
Date: Wed, 5 Aug 2020 14:03:34 +0200
Subject: [PATCH 04/19] add test to ensure LoadTestShape is not in user_classes

---
 locust/test/mock_locustfile.py | 4 +++-
 locust/test/test_main.py       | 2 ++
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/locust/test/mock_locustfile.py b/locust/test/mock_locustfile.py
index 8a80af31db..601f2a15d8 100644
--- a/locust/test/mock_locustfile.py
+++ b/locust/test/mock_locustfile.py
@@ -8,7 +8,7 @@
 MOCK_LOUCSTFILE_CONTENT = '''
 """This is a mock locust file for unit testing"""
 
-from locust import HttpUser, TaskSet, task, between
+from locust import HttpUser, TaskSet, task, between, LoadTestShape
 
 
 def index(l):
@@ -32,6 +32,8 @@ class UserSubclass(HttpUser):
 class NotUserSubclass():
     host = "http://localhost:8000"
 
+class LoadTestShape(LoadTestShape):
+    pass
 '''
 
 class MockedLocustfile:
diff --git a/locust/test/test_main.py b/locust/test/test_main.py
index 3a8aed84d9..ec0acff591 100644
--- a/locust/test/test_main.py
+++ b/locust/test/test_main.py
@@ -47,6 +47,7 @@ def test_load_locust_file_from_absolute_path(self):
             docstring, user_classes, shape_class = main.load_locustfile(mocked.file_path)
             self.assertIn('UserSubclass', user_classes)
             self.assertNotIn('NotUserSubclass', user_classes)
+            self.assertNotIn('LoadTestShape', user_classes)
 
     def test_load_locust_file_from_relative_path(self):
         with mock_locustfile() as mocked:
@@ -62,6 +63,7 @@ def test_return_docstring_and_user_classes(self):
             self.assertEqual("This is a mock locust file for unit testing", docstring)
             self.assertIn('UserSubclass', user_classes)
             self.assertNotIn('NotUserSubclass', user_classes)
+            self.assertNotIn('LoadTestShape', user_classes)
     
     def test_create_environment(self):
         options = parse_options(args=[

From 428c67fa44c69c25b92fe8685549ced6f664f1ab Mon Sep 17 00:00:00 2001
From: Max Williams <max.williams@deliveryhero.com>
Date: Wed, 5 Aug 2020 15:59:56 +0200
Subject: [PATCH 05/19] add a test

---
 locust/shape.py             |  6 +++++
 locust/test/test_runners.py | 51 ++++++++++++++++++++++++++++++++++++-
 2 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/locust/shape.py b/locust/shape.py
index efcea15ca0..4d50457658 100644
--- a/locust/shape.py
+++ b/locust/shape.py
@@ -9,6 +9,12 @@ class LoadTestShape(object):
 
     start_time = time.monotonic()
 
+    def reset_time(self):
+        '''
+        Resets start time back to 0
+        '''
+        self.start_time = time.monotonic()
+
     def get_run_time(self):
         '''
         Calculates run time in seconds of the load test
diff --git a/locust/test/test_runners.py b/locust/test/test_runners.py
index b24cc7e8e3..306ae9c8b8 100644
--- a/locust/test/test_runners.py
+++ b/locust/test/test_runners.py
@@ -6,7 +6,7 @@
 from gevent.queue import Queue
 
 import locust
-from locust import runners, between, constant
+from locust import runners, between, constant, LoadTestShape
 from locust.main import create_environment
 from locust.user import User, TaskSet, task
 from locust.env import Environment
@@ -812,6 +812,55 @@ def test_spawn_fewer_locusts_than_workers(self):
             
             self.assertEqual(2, num_users, "Total number of locusts that would have been spawned is not 2")
 
+    def test_custom_shape(self):
+        class MyUser(User):
+            wait_time = constant(0)
+            @task
+            def my_task(self):
+                pass
+    
+        class TestShape(LoadTestShape):
+            def tick(self):
+                run_time = self.get_run_time()
+                if run_time < 2:
+                    return (1, 1, False)
+                elif run_time < 4:
+                    return (2, 2, False)
+                else:
+                    return (0, 0, True)
+    
+        self.environment.user_classes = [MyUser]
+        self.environment.shape_class = TestShape()
+    
+        with mock.patch("locust.rpc.rpc.Server", mocked_rpc()) as server:
+            master = self.get_runner()
+            for i in range(5):
+                server.mocked_send(Message("client_ready", None, "fake_client%i" % i))
+    
+            # Start the shape_worker
+            self.environment.shape_class.reset_time()
+            master.start_shape()
+            sleep(0.5)
+    
+            # Wait for shape_worker to update user_count
+            num_users = 0
+            for _, msg in server.outbox:
+                if msg.data:
+                    num_users += msg.data["num_users"]
+            self.assertEqual(1, num_users, "Total number of users in first stage of shape test is not 1: %i" % num_users)
+    
+            # Wait for shape_worker to update user_count again
+            sleep(2)
+            num_users = 0
+            for _, msg in server.outbox:
+                if msg.data:
+                    num_users += msg.data["num_users"]
+            self.assertEqual(3, num_users, "Total number of users in second stage of shape test is not 3: %i" % num_users)
+            
+            # Wait to ensure shape_worker has stopped the test
+            sleep(3)
+            self.assertEqual('stopped', master.state, "The test has not been stopped by the shape class")
+
     def test_spawn_locusts_in_stepload_mode(self):
         with mock.patch("locust.rpc.rpc.Server", mocked_rpc()) as server:
             master = self.get_runner()

From c53b6256db725917d812e8514b785ea41d05c08b Mon Sep 17 00:00:00 2001
From: Max Williams <max.williams@deliveryhero.com>
Date: Wed, 5 Aug 2020 17:14:27 +0200
Subject: [PATCH 06/19] exit with error if a conflicting argument is specified

---
 locust/main.py    | 4 ++++
 locust/runners.py | 1 -
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/locust/main.py b/locust/main.py
index 04526b5d70..5b8f5e5f20 100644
--- a/locust/main.py
+++ b/locust/main.py
@@ -186,6 +186,10 @@ def main():
     # create locust Environment
     environment = create_environment(user_classes, options, events=locust.events, shape_class=shape_class)
 
+    if shape_class and (options.num_users or options.hatch_rate or options.run_time or options.step_load):
+        logger.error("The specified locustfile contains a shaper class but a conflicting argument was specified: users, hatch-rate, run-time or step-load")
+        sys.exit(1)
+
     if options.show_task_ratio:
         print("\n Task ratio per User class")
         print( "-" * 80)
diff --git a/locust/runners.py b/locust/runners.py
index 8933a3b1e6..8f5aa79561 100644
--- a/locust/runners.py
+++ b/locust/runners.py
@@ -309,7 +309,6 @@ def start_shape(self):
         self.shape_greenlet = self.greenlet.spawn(self.shape_worker)
         self.shape_greenlet.link_exception(greenlet_exception_handler)
 
-
     def shape_worker(self):
         logger.info('Shape worker starting')
         while self.state == STATE_INIT or self.state == STATE_HATCHING or self.state == STATE_RUNNING:

From 956a62973e66f223af4a7a9c423f2bba6106c1a9 Mon Sep 17 00:00:00 2001
From: Max Williams <max.williams@deliveryhero.com>
Date: Thu, 6 Aug 2020 15:05:34 +0200
Subject: [PATCH 07/19] Make formatting and style consistent

---
 examples/custom_shape/double_wave.py | 28 +++++++++++-----------------
 examples/custom_shape/stages.py      | 27 ++++++++++++++-------------
 examples/custom_shape/step_load.py   | 17 +++++++++--------
 locust/main.py                       |  4 ++--
 locust/runners.py                    |  8 ++++----
 locust/shape.py                      | 23 +++++++++++------------
 locust/test/test_runners.py          |  2 +-
 7 files changed, 52 insertions(+), 57 deletions(-)

diff --git a/examples/custom_shape/double_wave.py b/examples/custom_shape/double_wave.py
index 72ede36561..2f162ec410 100644
--- a/examples/custom_shape/double_wave.py
+++ b/examples/custom_shape/double_wave.py
@@ -15,25 +15,19 @@ class WebsiteUser(HttpUser):
 
 
 class DoubleWave(LoadTestShape):
-    '''
+    """
     A shape to immitate some specific user behaviour. In this example, midday
     and evening meal times.
 
-    Arguments:
-
-        min_users - minimum users
-        peak_one - first peak size
-        peak_two - second peak size
-        amplitude - size of the test
-        time_limit - total length of test
-    '''
-    def __init__(self,
-            min_users=20,
-            peak_one=1,
-            peak_two=1.5,
-            amplitude=150,
-            time_limit=600,
-        ):
+    Keyword arguments:
+
+        min_users -- minimum users
+        peak_one -- first peak size
+        peak_two -- second peak size
+        amplitude -- size of the test
+        time_limit -- total length of test
+    """
+    def __init__(self, min_users=20, peak_one=1, peak_two=1.5, amplitude=150, time_limit=600):
         self.min_users = min_users
         self.peak_one = peak_one
         self.peak_two = peak_two
@@ -49,7 +43,7 @@ def tick(self):
         run_time = self.get_run_time()
 
         if run_time < self.time_limit:
-            user_count = self.second_peak_users * math.e** - ((run_time/(math.sqrt(self.time_limit)*2))-4)**2 + self.first_peak_users * math.e** - ((run_time/(math.sqrt(self.time_limit)*2))-8)**2 + self.min_users
+            user_count = self.second_peak_users * math.e ** -((run_time/(math.sqrt(self.time_limit)*2))-4) ** 2 + self.first_peak_users * math.e ** - ((run_time/(math.sqrt(self.time_limit)*2))-8) ** 2 + self.min_users
             return (round(user_count), round(user_count), False)
         else:
             return (0, 0, True)
diff --git a/examples/custom_shape/stages.py b/examples/custom_shape/stages.py
index f4fc4216e8..8b10fcf382 100644
--- a/examples/custom_shape/stages.py
+++ b/examples/custom_shape/stages.py
@@ -14,21 +14,22 @@ class WebsiteUser(HttpUser):
 
 
 class StagesShape(LoadTestShape):
-    '''
+    """
     A simply load test shape class that has different user and hatch_rate at
     different stages.
 
-        duration - When this many seconds pass the test is advanced to the next stage
-        users - Total user count
-        hatch_rate - Hatch rate
-        stop - A boolean that can stop that test at a specific stage
+    Keyword arguments:
 
-    Arguments:
+        stages -- A list of dicts, each representing a stage with the following keys:
+            duration -- When this many seconds pass the test is advanced to the next stage
+            users -- Total user count
+            hatch_rate -- Hatch rate
+            stop -- A boolean that can stop that test at a specific stage
 
-        stop_at_end - can be set to stop once all stages have run.
-    '''
+        stop_at_end -- Can be set to stop once all stages have run.
+    """
     def __init__(self,
-            stages = [
+            stages=[
                 {'duration': 60, 'users': 10, 'hatch_rate': 10},
                 {'duration': 80, 'users': 30, 'hatch_rate': 10},
                 {'duration': 100, 'users': 50, 'hatch_rate': 10},
@@ -39,16 +40,16 @@ def __init__(self,
                 {'duration': 240, 'users': 1, 'hatch_rate': 1},
             ],
             stop_at_end=True
-        ):
-        self.stages = sorted(stages, key=lambda k: k['duration'])
+            ):
+        self.stages = sorted(stages, key=lambda k: k["duration"])
         self.stop_at_end = stop_at_end
 
     def tick(self):
         run_time = self.get_run_time()
 
         for stage in self.stages:
-            if run_time < stage['duration']:
-                tick_data = (stage['users'], stage['hatch_rate'], stage.get('stop', False))
+            if run_time < stage["duration"]:
+                tick_data = (stage["users"], stage["hatch_rate"], stage.get('stop', False))
                 self.last_stage = tick_data
                 return tick_data
 
diff --git a/examples/custom_shape/step_load.py b/examples/custom_shape/step_load.py
index 51ca14a8d4..b9615ee772 100644
--- a/examples/custom_shape/step_load.py
+++ b/examples/custom_shape/step_load.py
@@ -15,22 +15,23 @@ class WebsiteUser(HttpUser):
 
 
 class StepLoadShape(LoadTestShape):
-    '''
+    """
     A step load shape
 
-    Arguments:
 
-        step_time - Time between steps
-        step_load - User increase amount at each step
-        hatch_rate - Hatch rate to use at every step
-        time_limit - Time limit in seconds
-    '''
+    Keyword arguments:
+
+        step_time -- Time between steps
+        step_load -- User increase amount at each step
+        hatch_rate -- Hatch rate to use at every step
+        time_limit -- Time limit in seconds
+    """
     def __init__(self,
             step_time=30,
             step_load=10,
             hatch_rate=10,
             time_limit=600
-        ):
+            ):
         self.step_time = step_time
         self.step_load = step_load
         self.hatch_rate = hatch_rate
diff --git a/locust/main.py b/locust/main.py
index 5b8f5e5f20..5297e4c95f 100644
--- a/locust/main.py
+++ b/locust/main.py
@@ -186,8 +186,8 @@ def main():
     # create locust Environment
     environment = create_environment(user_classes, options, events=locust.events, shape_class=shape_class)
 
-    if shape_class and (options.num_users or options.hatch_rate or options.run_time or options.step_load):
-        logger.error("The specified locustfile contains a shaper class but a conflicting argument was specified: users, hatch-rate, run-time or step-load")
+    if shape_class and (options.num_users or options.hatch_rate or options.step_load):
+        logger.error("The specified locustfile contains a shaper class but a conflicting argument was specified: users, hatch-rate or step-load")
         sys.exit(1)
 
     if options.show_task_ratio:
diff --git a/locust/runners.py b/locust/runners.py
index 8f5aa79561..44e393c5e2 100644
--- a/locust/runners.py
+++ b/locust/runners.py
@@ -293,10 +293,10 @@ def stepload_worker(self, hatch_rate, step_users_growth, step_duration):
         while self.state == STATE_INIT or self.state == STATE_HATCHING or self.state == STATE_RUNNING:
             current_num_users += step_users_growth
             if current_num_users > int(self.total_users):
-                logger.info('Step Load is finished.')
+                logger.info("Step Load is finished")
                 break
             self.start(current_num_users, hatch_rate)
-            logger.info('Step loading: start hatch job of %d user.' % (current_num_users))
+            logger.info("Step loading: start hatch job of %d user" % (current_num_users))
             gevent.sleep(step_duration)
 
     def start_shape(self):
@@ -310,12 +310,12 @@ def start_shape(self):
         self.shape_greenlet.link_exception(greenlet_exception_handler)
 
     def shape_worker(self):
-        logger.info('Shape worker starting')
+        logger.info("Shape worker starting")
         while self.state == STATE_INIT or self.state == STATE_HATCHING or self.state == STATE_RUNNING:
             new_state = self.environment.shape_class.tick()
             user_count, hatch_rate, stop_test = new_state
             if stop_test:
-                logger.info('Shape test stopping')
+                logger.info("Shape test stopping")
                 self.stop()
             elif self.shape_last_state == new_state:
                 gevent.sleep(1)
diff --git a/locust/shape.py b/locust/shape.py
index 4d50457658..f82c1b2264 100644
--- a/locust/shape.py
+++ b/locust/shape.py
@@ -2,35 +2,34 @@
 
 
 class LoadTestShape(object):
-    '''
+    """
     A simple load test shape class used to control the shape of load generated
     during a load test.
-    '''
+    """
 
     start_time = time.monotonic()
 
     def reset_time(self):
-        '''
+        """
         Resets start time back to 0
-        '''
+        """
         self.start_time = time.monotonic()
 
     def get_run_time(self):
-        '''
+        """
         Calculates run time in seconds of the load test
-        '''
+        """
         return round(time.monotonic() - self.start_time)
-        
 
     def tick(self):
-        '''
+        """
         Returns a tuple with 3 elements to control the running load test.
 
-            user_count: Total user count
-            hatch_rate: Hatch rate to use when changing total user count
-            stop_test: A boolean to stop the test
+            user_count -- Total user count
+            hatch_rate -- Hatch rate to use when changing total user count
+            stop_test -- A boolean to stop the test
 
-        '''
+        """
         run_time = self.get_run_time()
 
         return (0, 0, True)
diff --git a/locust/test/test_runners.py b/locust/test/test_runners.py
index 306ae9c8b8..812451af8d 100644
--- a/locust/test/test_runners.py
+++ b/locust/test/test_runners.py
@@ -859,7 +859,7 @@ def tick(self):
             
             # Wait to ensure shape_worker has stopped the test
             sleep(3)
-            self.assertEqual('stopped', master.state, "The test has not been stopped by the shape class")
+            self.assertEqual("stopped", master.state, "The test has not been stopped by the shape class")
 
     def test_spawn_locusts_in_stepload_mode(self):
         with mock.patch("locust.rpc.rpc.Server", mocked_rpc()) as server:

From efcef244de408523c4dd852d08d17dcb359762e5 Mon Sep 17 00:00:00 2001
From: Max Williams <max.williams@deliveryhero.com>
Date: Fri, 7 Aug 2020 11:00:38 +0200
Subject: [PATCH 08/19] Remove double log

---
 locust/runners.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/locust/runners.py b/locust/runners.py
index 44e393c5e2..cf59573d69 100644
--- a/locust/runners.py
+++ b/locust/runners.py
@@ -77,9 +77,6 @@ def on_hatch_complete(user_count):
                 self.stats.reset_all()
         self.environment.events.hatch_complete.add_listener(on_hatch_complete)
 
-        if self.environment.shape_class:
-            logger.debug("Starting with shape class")
-
     def __del__(self):
         # don't leave any stray greenlets if runner is removed
         if self.greenlet and len(self.greenlet) > 0:

From bf932ae745b55514ee2541484790effd44013721 Mon Sep 17 00:00:00 2001
From: Max Williams <max.williams@deliveryhero.com>
Date: Fri, 7 Aug 2020 11:18:47 +0200
Subject: [PATCH 09/19] Add test for scaling down also

---
 locust/test/test_runners.py | 51 ++++++++++++++++++++++++++++++++++++-
 1 file changed, 50 insertions(+), 1 deletion(-)

diff --git a/locust/test/test_runners.py b/locust/test/test_runners.py
index 812451af8d..08239b483a 100644
--- a/locust/test/test_runners.py
+++ b/locust/test/test_runners.py
@@ -812,7 +812,7 @@ def test_spawn_fewer_locusts_than_workers(self):
             
             self.assertEqual(2, num_users, "Total number of locusts that would have been spawned is not 2")
 
-    def test_custom_shape(self):
+    def test_custom_shape_scale_up(self):
         class MyUser(User):
             wait_time = constant(0)
             @task
@@ -861,6 +861,55 @@ def tick(self):
             sleep(3)
             self.assertEqual("stopped", master.state, "The test has not been stopped by the shape class")
 
+    def test_custom_shape_scale_down(self):
+        class MyUser(User):
+            wait_time = constant(0)
+            @task
+            def my_task(self):
+                pass
+    
+        class TestShape(LoadTestShape):
+            def tick(self):
+                run_time = self.get_run_time()
+                if run_time < 2:
+                    return (5, 5, False)
+                elif run_time < 4:
+                    return (-4, 4, False)
+                else:
+                    return (0, 0, True)
+    
+        self.environment.user_classes = [MyUser]
+        self.environment.shape_class = TestShape()
+    
+        with mock.patch("locust.rpc.rpc.Server", mocked_rpc()) as server:
+            master = self.get_runner()
+            for i in range(5):
+                server.mocked_send(Message("client_ready", None, "fake_client%i" % i))
+    
+            # Start the shape_worker
+            self.environment.shape_class.reset_time()
+            master.start_shape()
+            sleep(0.5)
+    
+            # Wait for shape_worker to update user_count
+            num_users = 0
+            for _, msg in server.outbox:
+                if msg.data:
+                    num_users += msg.data["num_users"]
+            self.assertEqual(5, num_users, "Total number of users in first stage of shape test is not 5: %i" % num_users)
+    
+            # Wait for shape_worker to update user_count again
+            sleep(2)
+            num_users = 0
+            for _, msg in server.outbox:
+                if msg.data:
+                    num_users += msg.data["num_users"]
+            self.assertEqual(1, num_users, "Total number of users in second stage of shape test is not 1: %i" % num_users)
+            
+            # Wait to ensure shape_worker has stopped the test
+            sleep(3)
+            self.assertEqual("stopped", master.state, "The test has not been stopped by the shape class")
+
     def test_spawn_locusts_in_stepload_mode(self):
         with mock.patch("locust.rpc.rpc.Server", mocked_rpc()) as server:
             master = self.get_runner()

From f593eb273e1ac110e72bc8b7e466a23ebd3c230d Mon Sep 17 00:00:00 2001
From: Max Williams <max.williams@deliveryhero.com>
Date: Fri, 7 Aug 2020 11:21:24 +0200
Subject: [PATCH 10/19] Revert removing \n

---
 locust/runners.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/locust/runners.py b/locust/runners.py
index cf59573d69..b1df0871fe 100644
--- a/locust/runners.py
+++ b/locust/runners.py
@@ -73,7 +73,7 @@ def on_request_failure(request_type, name, response_time, response_length, excep
         def on_hatch_complete(user_count):
             self.state = STATE_RUNNING
             if environment.reset_stats:
-                logger.info("Resetting stats")
+                logger.info("Resetting stats\n")
                 self.stats.reset_all()
         self.environment.events.hatch_complete.add_listener(on_hatch_complete)
 

From 1e20f8851dcc88a2ab013161c806bd551d0df5cb Mon Sep 17 00:00:00 2001
From: Max Williams <max.williams@deliveryhero.com>
Date: Mon, 10 Aug 2020 14:09:09 +0200
Subject: [PATCH 11/19] rewrite DoubleWave example

---
 examples/custom_shape/double_wave.py | 26 +++++++++++++-------------
 1 file changed, 13 insertions(+), 13 deletions(-)

diff --git a/examples/custom_shape/double_wave.py b/examples/custom_shape/double_wave.py
index 2f162ec410..9e2cec277f 100644
--- a/examples/custom_shape/double_wave.py
+++ b/examples/custom_shape/double_wave.py
@@ -19,31 +19,31 @@ class DoubleWave(LoadTestShape):
     A shape to immitate some specific user behaviour. In this example, midday
     and evening meal times.
 
-    Keyword arguments:
+    Settings:
 
         min_users -- minimum users
-        peak_one -- first peak size
-        peak_two -- second peak size
-        amplitude -- size of the test
+        peak_one_users -- users in first peak
+        peak_two_users -- users in second peak
         time_limit -- total length of test
+
     """
-    def __init__(self, min_users=20, peak_one=1, peak_two=1.5, amplitude=150, time_limit=600):
-        self.min_users = min_users
-        self.peak_one = peak_one
-        self.peak_two = peak_two
-        self.amplitude = amplitude
-        self.time_limit = time_limit
+    min_users = 20
+    peak_one_users = 60
+    peak_two_users = 40
+    time_limit = 600
 
+    def __init__(self):
         self.first_peak_time = round(self.time_limit / 3)
-        self.first_peak_users = round(self.peak_one * self.amplitude)
         self.second_peak_time = round(self.first_peak_time * 2)
-        self.second_peak_users = round(self.peak_two * self.amplitude)
 
     def tick(self):
         run_time = self.get_run_time()
 
         if run_time < self.time_limit:
-            user_count = self.second_peak_users * math.e ** -((run_time/(math.sqrt(self.time_limit)*2))-4) ** 2 + self.first_peak_users * math.e ** - ((run_time/(math.sqrt(self.time_limit)*2))-8) ** 2 + self.min_users
+            user_count = (self.peak_one_users - self.min_users) * math.e ** -((run_time/(math.sqrt(self.time_limit)*2))-4) ** 2
+            + (self.peak_two_users - self.min_users) * math.e ** -((run_time/(math.sqrt(self.time_limit)*2))-8) ** 2
+            + self.min_users
+
             return (round(user_count), round(user_count), False)
         else:
             return (0, 0, True)

From d84d0c084e2a3c3a9d74f59d7e7a4212b42c0da5 Mon Sep 17 00:00:00 2001
From: Max Williams <max.williams@deliveryhero.com>
Date: Mon, 10 Aug 2020 14:22:32 +0200
Subject: [PATCH 12/19] simplify stages example

---
 examples/custom_shape/stages.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/examples/custom_shape/stages.py b/examples/custom_shape/stages.py
index 8b10fcf382..794847ff7c 100644
--- a/examples/custom_shape/stages.py
+++ b/examples/custom_shape/stages.py
@@ -31,12 +31,10 @@ class StagesShape(LoadTestShape):
     def __init__(self,
             stages=[
                 {'duration': 60, 'users': 10, 'hatch_rate': 10},
-                {'duration': 80, 'users': 30, 'hatch_rate': 10},
                 {'duration': 100, 'users': 50, 'hatch_rate': 10},
-                {'duration': 120, 'users': 70, 'hatch_rate': 10},
                 {'duration': 180, 'users': 100, 'hatch_rate': 10},
-                {'duration': 220, 'users': 10, 'hatch_rate': 10},
-                {'duration': 230, 'users': 5, 'hatch_rate': 10},
+                {'duration': 220, 'users': 30, 'hatch_rate': 10},
+                {'duration': 230, 'users': 10, 'hatch_rate': 10},
                 {'duration': 240, 'users': 1, 'hatch_rate': 1},
             ],
             stop_at_end=True

From 37271ecbf350ce0040baef78ccc8e99c7d2d26f9 Mon Sep 17 00:00:00 2001
From: Max Williams <max.williams@deliveryhero.com>
Date: Mon, 10 Aug 2020 14:24:44 +0200
Subject: [PATCH 13/19] y u round, no round

---
 locust/shape.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/locust/shape.py b/locust/shape.py
index f82c1b2264..c23a3aaa3a 100644
--- a/locust/shape.py
+++ b/locust/shape.py
@@ -19,7 +19,7 @@ def get_run_time(self):
         """
         Calculates run time in seconds of the load test
         """
-        return round(time.monotonic() - self.start_time)
+        return time.monotonic() - self.start_time
 
     def tick(self):
         """

From f921aea1f01b90127664949f144a114274b82c76 Mon Sep 17 00:00:00 2001
From: Max Williams <max.williams@deliveryhero.com>
Date: Tue, 11 Aug 2020 17:53:46 +0200
Subject: [PATCH 14/19] Simplify examples, return None to stop test

---
 examples/custom_shape/double_wave.py | 20 +++++++------------
 examples/custom_shape/stages.py      | 30 ++++++++++------------------
 examples/custom_shape/step_load.py   | 22 ++++++++++----------
 locust/runners.py                    |  6 +++---
 locust/shape.py                      |  7 ++++---
 locust/test/test_runners.py          | 12 +++++------
 6 files changed, 42 insertions(+), 55 deletions(-)

diff --git a/examples/custom_shape/double_wave.py b/examples/custom_shape/double_wave.py
index 9e2cec277f..d1b26ea121 100644
--- a/examples/custom_shape/double_wave.py
+++ b/examples/custom_shape/double_wave.py
@@ -20,30 +20,24 @@ class DoubleWave(LoadTestShape):
     and evening meal times.
 
     Settings:
-
         min_users -- minimum users
         peak_one_users -- users in first peak
         peak_two_users -- users in second peak
         time_limit -- total length of test
-
     """
+
     min_users = 20
     peak_one_users = 60
     peak_two_users = 40
     time_limit = 600
 
-    def __init__(self):
-        self.first_peak_time = round(self.time_limit / 3)
-        self.second_peak_time = round(self.first_peak_time * 2)
-
     def tick(self):
-        run_time = self.get_run_time()
+        run_time = round(self.get_run_time())
 
         if run_time < self.time_limit:
-            user_count = (self.peak_one_users - self.min_users) * math.e ** -((run_time/(math.sqrt(self.time_limit)*2))-4) ** 2
-            + (self.peak_two_users - self.min_users) * math.e ** -((run_time/(math.sqrt(self.time_limit)*2))-8) ** 2
-            + self.min_users
-
-            return (round(user_count), round(user_count), False)
+            user_count = ((self.peak_one_users - self.min_users) * math.e ** -((run_time/(self.time_limit/10*2/3))-5) ** 2
+                        + (self.peak_two_users - self.min_users) * math.e ** - ((run_time/(self.time_limit/10*2/3))-10) ** 2
+                        + self.min_users)
+            return (round(user_count), round(user_count))
         else:
-            return (0, 0, True)
+            return None
diff --git a/examples/custom_shape/stages.py b/examples/custom_shape/stages.py
index 794847ff7c..83346cfe10 100644
--- a/examples/custom_shape/stages.py
+++ b/examples/custom_shape/stages.py
@@ -28,30 +28,22 @@ class StagesShape(LoadTestShape):
 
         stop_at_end -- Can be set to stop once all stages have run.
     """
-    def __init__(self,
-            stages=[
-                {'duration': 60, 'users': 10, 'hatch_rate': 10},
-                {'duration': 100, 'users': 50, 'hatch_rate': 10},
-                {'duration': 180, 'users': 100, 'hatch_rate': 10},
-                {'duration': 220, 'users': 30, 'hatch_rate': 10},
-                {'duration': 230, 'users': 10, 'hatch_rate': 10},
-                {'duration': 240, 'users': 1, 'hatch_rate': 1},
-            ],
-            stop_at_end=True
-            ):
-        self.stages = sorted(stages, key=lambda k: k["duration"])
-        self.stop_at_end = stop_at_end
+
+    stages = [
+        {'duration': 60, 'users': 10, 'hatch_rate': 10},
+        {'duration': 100, 'users': 50, 'hatch_rate': 10},
+        {'duration': 180, 'users': 100, 'hatch_rate': 10},
+        {'duration': 220, 'users': 30, 'hatch_rate': 10},
+        {'duration': 230, 'users': 10, 'hatch_rate': 10},
+        {'duration': 240, 'users': 1, 'hatch_rate': 1},
+    ]
 
     def tick(self):
         run_time = self.get_run_time()
 
         for stage in self.stages:
             if run_time < stage["duration"]:
-                tick_data = (stage["users"], stage["hatch_rate"], stage.get('stop', False))
-                self.last_stage = tick_data
+                tick_data = (stage["users"], stage["hatch_rate"])
                 return tick_data
 
-        if self.stop_at_end:
-            return (0, 0, True)
-        else:
-            return self.last_stage
+        return None
diff --git a/examples/custom_shape/step_load.py b/examples/custom_shape/step_load.py
index b9615ee772..a8658586f1 100644
--- a/examples/custom_shape/step_load.py
+++ b/examples/custom_shape/step_load.py
@@ -25,19 +25,19 @@ class StepLoadShape(LoadTestShape):
         step_load -- User increase amount at each step
         hatch_rate -- Hatch rate to use at every step
         time_limit -- Time limit in seconds
+
     """
-    def __init__(self,
-            step_time=30,
-            step_load=10,
-            hatch_rate=10,
-            time_limit=600
-            ):
-        self.step_time = step_time
-        self.step_load = step_load
-        self.hatch_rate = hatch_rate
-        self.time_limit = time_limit
+
+    step_time = 30
+    step_load = 10
+    hatch_rate = 10
+    time_limit = 600
 
     def tick(self):
         run_time = self.get_run_time()
+
+        if run_time > self.time_limit:
+            return None
+
         current_step = math.floor(run_time / self.step_time) + 1
-        return (current_step * self.step_load, self.hatch_rate, run_time > self.time_limit)
+        return (current_step * self.step_load, self.hatch_rate)
diff --git a/locust/runners.py b/locust/runners.py
index b1df0871fe..2343ceb076 100644
--- a/locust/runners.py
+++ b/locust/runners.py
@@ -76,7 +76,7 @@ def on_hatch_complete(user_count):
                 logger.info("Resetting stats\n")
                 self.stats.reset_all()
         self.environment.events.hatch_complete.add_listener(on_hatch_complete)
-
+    
     def __del__(self):
         # don't leave any stray greenlets if runner is removed
         if self.greenlet and len(self.greenlet) > 0:
@@ -310,13 +310,13 @@ def shape_worker(self):
         logger.info("Shape worker starting")
         while self.state == STATE_INIT or self.state == STATE_HATCHING or self.state == STATE_RUNNING:
             new_state = self.environment.shape_class.tick()
-            user_count, hatch_rate, stop_test = new_state
-            if stop_test:
+            if new_state == None:
                 logger.info("Shape test stopping")
                 self.stop()
             elif self.shape_last_state == new_state:
                 gevent.sleep(1)
             else:
+                user_count, hatch_rate = new_state
                 logger.info("Shape test updating to %d users at %.2f hatch rate" % (user_count, hatch_rate))
                 self.start(user_count=user_count, hatch_rate=hatch_rate)
                 self.shape_last_state = new_state
diff --git a/locust/shape.py b/locust/shape.py
index c23a3aaa3a..3ef392781b 100644
--- a/locust/shape.py
+++ b/locust/shape.py
@@ -23,13 +23,14 @@ def get_run_time(self):
 
     def tick(self):
         """
-        Returns a tuple with 3 elements to control the running load test.
+        Returns a tuple with 2 elements to control the running load test:
 
             user_count -- Total user count
             hatch_rate -- Hatch rate to use when changing total user count
-            stop_test -- A boolean to stop the test
+
+        If `None` is returned then the running load test will be stopped.
 
         """
         run_time = self.get_run_time()
 
-        return (0, 0, True)
+        return None
diff --git a/locust/test/test_runners.py b/locust/test/test_runners.py
index 08239b483a..6113b32b2b 100644
--- a/locust/test/test_runners.py
+++ b/locust/test/test_runners.py
@@ -823,11 +823,11 @@ class TestShape(LoadTestShape):
             def tick(self):
                 run_time = self.get_run_time()
                 if run_time < 2:
-                    return (1, 1, False)
+                    return (1, 1)
                 elif run_time < 4:
-                    return (2, 2, False)
+                    return (2, 2)
                 else:
-                    return (0, 0, True)
+                    return None
     
         self.environment.user_classes = [MyUser]
         self.environment.shape_class = TestShape()
@@ -872,11 +872,11 @@ class TestShape(LoadTestShape):
             def tick(self):
                 run_time = self.get_run_time()
                 if run_time < 2:
-                    return (5, 5, False)
+                    return (5, 5)
                 elif run_time < 4:
-                    return (-4, 4, False)
+                    return (-4, 4)
                 else:
-                    return (0, 0, True)
+                    return None
     
         self.environment.user_classes = [MyUser]
         self.environment.shape_class = TestShape()

From e32ec318535c8396c4d00da87cefd57c2ae63319 Mon Sep 17 00:00:00 2001
From: Max Williams <max.williams@deliveryhero.com>
Date: Wed, 12 Aug 2020 17:57:58 +0200
Subject: [PATCH 15/19] is None? Yes is None.

---
 locust/runners.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/locust/runners.py b/locust/runners.py
index 2343ceb076..188a0c72ff 100644
--- a/locust/runners.py
+++ b/locust/runners.py
@@ -310,7 +310,7 @@ def shape_worker(self):
         logger.info("Shape worker starting")
         while self.state == STATE_INIT or self.state == STATE_HATCHING or self.state == STATE_RUNNING:
             new_state = self.environment.shape_class.tick()
-            if new_state == None:
+            if new_state is None:
                 logger.info("Shape test stopping")
                 self.stop()
             elif self.shape_last_state == new_state:

From 36738ad2c9a1f6fa437aeab218711440a38ad051 Mon Sep 17 00:00:00 2001
From: Max Williams <max.williams@deliveryhero.com>
Date: Fri, 14 Aug 2020 15:09:33 +0200
Subject: [PATCH 16/19] disable editing running test. Print message about
 user_count/hatch_rate being ignored

---
 locust/runners.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/locust/runners.py b/locust/runners.py
index 188a0c72ff..fa9067029d 100644
--- a/locust/runners.py
+++ b/locust/runners.py
@@ -297,11 +297,11 @@ def stepload_worker(self, hatch_rate, step_users_growth, step_duration):
             gevent.sleep(step_duration)
 
     def start_shape(self):
-        logger.info("Shape test starting")
         if self.shape_greenlet:
-            logger.info("There is an ongoing shape test running, will stop it now")
-            self.shape_greenlet.kill()
+            logger.info("There is an ongoing shape test running. Editing is disabled")
+            return
 
+        logger.info("Shape test starting. User count and hatch rate are ignored for this type of load test")
         self.state = STATE_INIT
         self.shape_greenlet = self.greenlet.spawn(self.shape_worker)
         self.shape_greenlet.link_exception(greenlet_exception_handler)

From ecad2f722913f2a3b17a6152b88d9b76a15c7b13 Mon Sep 17 00:00:00 2001
From: Max Williams <max.williams@deliveryhero.com>
Date: Fri, 14 Aug 2020 16:30:55 +0200
Subject: [PATCH 17/19] Adding doc

---
 docs/generating-custom-load-shape.rst | 38 +++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)
 create mode 100644 docs/generating-custom-load-shape.rst

diff --git a/docs/generating-custom-load-shape.rst b/docs/generating-custom-load-shape.rst
new file mode 100644
index 0000000000..8bfdbd2f86
--- /dev/null
+++ b/docs/generating-custom-load-shape.rst
@@ -0,0 +1,38 @@
+.. _generating-custom-load-shape:
+
+=================================
+Generating a custom load shape using a `LoadTestShape` class
+=================================
+
+Sometimes a completely custom shaped load test is required that cannot be achieved by simply setting or changing the user count and hatch rate. For example, generating a spike during a test or ramping up and down at custom times. In these cases using the `LoadTestShape` class can give complete control over the user count and hatch rate at all times.
+
+How does a `LoadTestShape` class work?
+---------------------------------------------
+
+Define your class inheriting the `LoadTestShape` class in your locust file. If this type of class is found then it will be automatically used by Locust. Add a `tick()` method to return either a tuple containing the user count and hatch rate or `None` to stop the load test. Locust will call the `tick()` method every second and update the load test with new settings or stop the test.
+
+Examples
+---------------------------------------------
+
+There are also some [examples on github](https://github.com/locustio/locust/tree/master/examples/custom_shape) including:
+
+- Generating a double wave shape
+- Time based stages like K6
+- Step load pattern like Visual Studio
+
+Here is a simple example that will increase user count in blocks of 100 then stop the load test after 10 minutes:
+
+```python
+class MyCustomShape(LoadTestShape):
+    time_limit = 600
+    hatch_rate = 20
+    
+    def tick(self):
+        run_time = self.get_run_time()
+
+        if run_time < self.time_limit:
+            user_count = round(run_time, -2)
+            return (user_count, hatch_rate)
+
+        return None
+```

From e3ec440aeddf5c0a0ca808674659409e756040c0 Mon Sep 17 00:00:00 2001
From: Max Williams <max.williams@deliveryhero.com>
Date: Fri, 14 Aug 2020 19:28:06 +0200
Subject: [PATCH 18/19] update docs, change one word in log message

---
 docs/index.rst                      | 1 +
 docs/running-locust-distributed.rst | 6 ++++++
 locust/main.py                      | 2 +-
 3 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/docs/index.rst b/docs/index.rst
index 32574ec0a4..be36afeaa6 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -52,6 +52,7 @@ Other functionalities
 .. toctree ::
     :maxdepth: 2
 
+    generating-custom-load-shape
     running-locust-in-step-load-mode
     retrieving-stats
     testing-other-systems
diff --git a/docs/running-locust-distributed.rst b/docs/running-locust-distributed.rst
index b819f2ffa7..3d2c7691a9 100644
--- a/docs/running-locust-distributed.rst
+++ b/docs/running-locust-distributed.rst
@@ -100,6 +100,12 @@ Running Locust distributed without the web UI
 See :ref:`running-locust-distributed-without-web-ui`
 
 
+Generating a custom load shape using a `LoadTestShape` class
+=============================================
+
+See :ref:`generating-custom-load-shape`
+
+
 Running Locust distributed in Step Load mode
 =============================================
 
diff --git a/locust/main.py b/locust/main.py
index 5297e4c95f..2fea2cbda5 100644
--- a/locust/main.py
+++ b/locust/main.py
@@ -187,7 +187,7 @@ def main():
     environment = create_environment(user_classes, options, events=locust.events, shape_class=shape_class)
 
     if shape_class and (options.num_users or options.hatch_rate or options.step_load):
-        logger.error("The specified locustfile contains a shaper class but a conflicting argument was specified: users, hatch-rate or step-load")
+        logger.error("The specified locustfile contains a shape class but a conflicting argument was specified: users, hatch-rate or step-load")
         sys.exit(1)
 
     if options.show_task_ratio:

From 93a41eead95b287e3021c084aed4117339af527e Mon Sep 17 00:00:00 2001
From: Max Williams <max.williams@deliveryhero.com>
Date: Mon, 17 Aug 2020 11:02:21 +0200
Subject: [PATCH 19/19] Add full test with workers and master

---
 locust/test/test_runners.py | 56 +++++++++++++++++++++++++++++++++++++
 1 file changed, 56 insertions(+)

diff --git a/locust/test/test_runners.py b/locust/test/test_runners.py
index 6113b32b2b..3848759462 100644
--- a/locust/test/test_runners.py
+++ b/locust/test/test_runners.py
@@ -416,6 +416,62 @@ def incr_stats(l):
             20, 
             "For some reason the master node's stats has not come in",
         )
+    
+    def test_distributed_shape(self):
+        """
+        Full integration test that starts both a MasterRunner and three WorkerRunner instances 
+        and tests a basic LoadTestShape with scaling up and down users
+        """
+        class TestUser(User):
+            wait_time = constant(0)
+            @task
+            def my_task(self):
+                pass
+    
+        class TestShape(LoadTestShape):
+            def tick(self):
+                run_time = self.get_run_time()
+                if run_time < 2:
+                    return (9, 9)
+                elif run_time < 4:
+                    return (21, 21)
+                elif run_time < 6:
+                    return (3, 21)
+                else:
+                    return None
+        
+        with mock.patch("locust.runners.WORKER_REPORT_INTERVAL", new=0.3):
+            master_env = Environment(user_classes=[TestUser], shape_class=TestShape())
+            master_env.shape_class.reset_time()
+            master = master_env.create_master_runner("*", 0)
+
+            workers = []
+            for i in range(3):
+                worker_env = Environment(user_classes=[TestUser])
+                worker = worker_env.create_worker_runner("127.0.0.1", master.server.port)
+                workers.append(worker)
+    
+            # Give workers time to connect
+            sleep(0.1)
+            # Start a shape test
+            master.start_shape()
+            sleep(1)
+
+            # Ensure workers have connected and started the correct amounf of users
+            for worker in workers:
+                self.assertEqual(3, worker.user_count, "Shape test has not reached stage 1")
+            # Ensure new stage with more users has been reached
+            sleep(2)
+            for worker in workers:
+                self.assertEqual(7, worker.user_count, "Shape test has not reached stage 2")
+            # Ensure new stage with less users has been reached
+            sleep(2)
+            for worker in workers:
+                self.assertEqual(1, worker.user_count, "Shape test has not reached stage 3")
+            # Ensure test stops at the end
+            sleep(2)
+            for worker in workers:
+                self.assertEqual(0, worker.user_count, "Shape test has not stopped")
 
 
 class TestMasterRunner(LocustTestCase):