From 7040ad1aae40028ef7b4390abd527a51241fea0e Mon Sep 17 00:00:00 2001
From: Maxence Boutet <mboutet@explorance.com>
Date: Mon, 27 Sep 2021 10:58:55 -0400
Subject: [PATCH] Ensure target_user_count is set before test_start event is
 fired

---
 locust/runners.py           |  9 +++++----
 locust/test/test_runners.py | 18 ++++++++++++++++++
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/locust/runners.py b/locust/runners.py
index 8805ee5294..f4afaa4ccb 100644
--- a/locust/runners.py
+++ b/locust/runners.py
@@ -303,6 +303,8 @@ def start(self, user_count: int, spawn_rate: float, wait: bool = False):
                      If False (the default), a greenlet that spawns the users will be
                      started and the call to this method will return immediately.
         """
+        self.target_user_count = user_count
+
         if self.state != STATE_RUNNING and self.state != STATE_SPAWNING:
             self.stats.clear_all()
             self.exceptions = {}
@@ -329,8 +331,6 @@ def start(self, user_count: int, spawn_rate: float, wait: bool = False):
 
         self._users_dispatcher.new_dispatch(user_count, spawn_rate)
 
-        self.target_user_count = user_count
-
         try:
             for dispatched_users in self._users_dispatcher:
                 user_classes_spawn_count = {}
@@ -659,6 +659,9 @@ def cpu_log_warning(self):
 
     def start(self, user_count: int, spawn_rate: float, **kwargs) -> None:
         self.spawning_completed = False
+
+        self.target_user_count = user_count
+
         num_workers = len(self.clients.ready) + len(self.clients.running) + len(self.clients.spawning)
         if not num_workers:
             logger.warning(
@@ -700,8 +703,6 @@ def start(self, user_count: int, spawn_rate: float, **kwargs) -> None:
 
         self._users_dispatcher.new_dispatch(target_user_count=user_count, spawn_rate=spawn_rate)
 
-        self.target_user_count = user_count
-
         try:
             for dispatched_users in self._users_dispatcher:
                 dispatch_greenlets = Group()
diff --git a/locust/test/test_runners.py b/locust/test/test_runners.py
index 69a6ca8859..8ac4a2518c 100644
--- a/locust/test/test_runners.py
+++ b/locust/test/test_runners.py
@@ -621,6 +621,13 @@ def my_task(self):
         environment = Environment(user_classes=[MyUser1])
         runner = LocalRunner(environment)
 
+        test_start_event_fired = [False]
+
+        @environment.events.test_start.add_listener
+        def on_test_start(*args, **kwargs):
+            test_start_event_fired[0] = True
+            self.assertEqual(runner.target_user_count, 3)
+
         runner.start(user_count=3, spawn_rate=1, wait=False)
 
         gevent.sleep(1)
@@ -639,6 +646,8 @@ def my_task(self):
 
         runner.quit()
 
+        self.assertTrue(test_start_event_fired[0])
+
 
 class TestMasterWorkerRunners(LocustTestCase):
     def test_distributed_integration_run(self):
@@ -1594,6 +1603,13 @@ def my_task(self):
             master_env = Environment(user_classes=[MyUser1])
             master = master_env.create_master_runner("*", 0)
 
+            test_start_event_fired = [False]
+
+            @master_env.events.test_start.add_listener
+            def on_test_start(*args, **kwargs):
+                test_start_event_fired[0] = True
+                self.assertEqual(master.target_user_count, 3)
+
             sleep(0)
 
             # start 1 worker runner
@@ -1624,6 +1640,8 @@ def my_task(self):
             # make sure users are killed
             self.assertEqual(0, worker.user_count)
 
+            self.assertTrue(test_start_event_fired[0])
+
 
 class TestMasterRunner(LocustTestCase):
     def setUp(self):