update test case for step load check

locust/runners.py

@@ -346,8 +346,7 @@ def start_stepload(self, locust_count, hatch_rate, step_locust_count, step_durat
         if self.stepload_greenlet:
             logger.info("There is an ongoing swarming in Step Load mode, will stop it now.")
             self.greenlet.killone(self.stepload_greenlet)
-
-        logger.info("Start a new swarming in Step Load mode: %d locusts, %d delta locusts in step, %ds step duration " % (locust_count, step_locust_count, step_duration))
+        logger.info("Start a new swarming in Step Load mode: total locust count of %d, hatch rate of %d, step locust count of %d, step duration of %d " % (locust_count, hatch_rate, step_locust_count, step_duration))
         self.state = STATE_INIT
         self.stepload_greenlet = self.greenlet.spawn(self.stepload_worker)
         self.stepload_greenlet.link_exception(callback=self.noop)

locust/test/test_runners.py

@@ -56,7 +56,7 @@ def __init__(self):
         self.master_bind_host = '*'
         self.master_bind_port = 5557
         self.heartbeat_liveness = 3
-        self.heartbeat_interval = 0.01
+        self.heartbeat_interval = 3
         self.stop_timeout = None
         self.step_load = True
 
@@ -420,16 +420,30 @@ class MyTestLocust(Locust):
             master = MasterLocustRunner(MyTestLocust, self.options)
             for i in range(5):
                 server.mocked_send(Message("client_ready", None, "fake_client%i" % i))
-
-            master.start_stepload(2, 1, 1, 3)
+
+            # start a new swarming in Step Load mode: total locust count of 10, hatch rate of 2, step locust count of 5, step duration of 5s
+            master.start_stepload(10, 2, 5, 5)
+
+            # make sure the first step run is started
             sleep(1)
             self.assertEqual(5, len(server.outbox))
-            
+
             num_clients = 0
+            end_of_last_step = len(server.outbox)
             for _, msg in server.outbox:
                 num_clients += Message.unserialize(msg).data["num_clients"]
 
-            self.assertEqual(1, num_clients, "Total number of locusts that would have been spawned is not 1")
+            self.assertEqual(5, num_clients, "Total number of locusts that would have been spawned for first step is not 5")
+
+            # make sure the first step run is complete
+            sleep(5)
+            num_clients = 0
+            idx = end_of_last_step
+            while idx < len(server.outbox):
+                msg = server.outbox[idx][1]
+                num_clients += Message.unserialize(msg).data["num_clients"]
+                idx += 1
+            self.assertEqual(10, num_clients, "Total number of locusts that would have been spawned for second step is not 10")
 
     def test_exception_in_task(self):
         class HeyAnException(Exception):