Fix test_start and test_stop issues with new code

locustio · Jul 12, 2021 · 59936ac · 59936ac
1 parent f85af23
commit 59936ac
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 4 deletions.
diff --git a/locust/runners.py b/locust/runners.py
@@ -489,7 +489,6 @@ def stop(self):
         if self.state == STATE_STOPPED:
             return
         super().stop()
-        self.environment.events.test_stop.fire(environment=self.environment)
 
     def send_message(self, msg_type, data=None):
         """
@@ -805,7 +804,6 @@ def stop(self, send_stop_to_client: bool = True):
                     logger.error("Timeout waiting for all workers to stop")
                 finally:
                     timeout.cancel()
-
             self.environment.events.test_stop.fire(environment=self.environment)
 
     def quit(self):
@@ -1060,12 +1058,14 @@ def start_worker(self, user_classes_count: Dict[str, int], **kwargs):
         :param user_classes_count: Users to run
         """
         self.target_user_classes_count = user_classes_count
-
         if self.worker_state != STATE_RUNNING and self.worker_state != STATE_SPAWNING:
             self.stats.clear_all()
             self.exceptions = {}
             self.cpu_warning_emitted = False
             self.worker_cpu_warning_emitted = False
+            self.environment.events.test_start.fire(environment=self.environment)
+
+        self.worker_state = STATE_SPAWNING
 
         for user_class in self.user_classes:
             if self.environment.host is not None:
@@ -1123,7 +1123,6 @@ def worker(self):
                 logger.error("RPCError found when receiving from master: %s" % (e))
                 continue
             if msg.type == "spawn":
-                self.worker_state = STATE_SPAWNING
                 self.client.send(Message("spawning", None, self.client_id))
                 job = msg.data
                 if job["timestamp"] <= last_received_spawn_timestamp:

diff --git a/locust/test/test_runners.py b/locust/test/test_runners.py
@@ -2707,6 +2707,8 @@ def on_test_start(*args, **kw):
                 Message(
                     "spawn",
                     {
+                        "timestamp": 1605538585,
+                        "user_classes_count": {"MyTestUser": 1},
                         "spawn_rate": 1,
                         "num_users": 1,
                         "host": "",
@@ -2730,6 +2732,8 @@ def on_test_start(*args, **kw):
                 Message(
                     "spawn",
                     {
+                        "timestamp": 1605538586,
+                        "user_classes_count": {"MyTestUser": 1},
                         "spawn_rate": 1,
                         "num_users": 1,
                         "host": "",
@@ -2746,6 +2750,8 @@ def on_test_start(*args, **kw):
                 Message(
                     "spawn",
                     {
+                        "timestamp": 1605538587,
+                        "user_classes_count": {"MyTestUser": 1},
                         "spawn_rate": 1,
                         "num_users": 1,
                         "host": "",
@@ -2784,6 +2790,8 @@ def on_test_stop(*args, **kw):
                 Message(
                     "spawn",
                     {
+                        "timestamp": 1605538585,
+                        "user_classes_count": {"MyTestUser": 1},
                         "spawn_rate": 1,
                         "num_users": 1,
                         "host": "",
@@ -2817,6 +2825,8 @@ def on_test_stop(*args, **kw):
                 Message(
                     "spawn",
                     {
+                        "timestamp": 1605538586,
+                        "user_classes_count": {"MyTestUser": 1},
                         "spawn_rate": 1,
                         "num_users": 1,
                         "host": "",