Repository: locustio/locust
File: locust/test/test_runners.py (Python)

Code before:
                self.assertEqual(
                    36, test_shape.get_current_user_count(), "Shape is not seeing stage 2 runner user count correctly"
                )
            self.assertDictEqual(
                master.reported_user_classes_count, {"FixedUser1": 1, "FixedUser2": 11, "TestUser": 24}
            )

            # Ensure new stage with less users has been reached
            # and expected count of the fixed users is present
            sleep(1)
            for worker in workers:
                self.assertEqual(4, worker.user_count, "Shape test has not reached stage 3")
                self.assertEqual(
                    12, test_shape.get_current_user_count(), "Shape is not seeing stage 3 runner user count correctly"
                )
            self.assertDictEqual(master.reported_user_classes_count, {"FixedUser1": 1, "FixedUser2": 11, "TestUser": 0})

            # Ensure test stops at the end
            sleep(0.5)
            for worker in workers:
                self.assertEqual(0, worker.user_count, "Shape test has not stopped")
                self.assertEqual(
                    0, test_shape.get_current_user_count(), "Shape is not seeing stopped runner user count correctly"
                )
            self.assertDictEqual(master.reported_user_classes_count, {"FixedUser1": 0, "FixedUser2": 0, "TestUser": 0})
            sleep(0.5)
            self.assertEqual(STATE_STOPPED, master.state)

    def test_distributed_shape_with_stop_timeout(self):
        """
        Full integration test that starts both a MasterRunner and five WorkerRunner instances
        and tests a basic LoadTestShape with scaling up and down users
        """

        class TestUser1(User):
            def start(self, group: Group):
                gevent.sleep(0.5)
                return super().start(group)

            @task
            def my_task(self):
                gevent.sleep(0)

        class TestUser2(User):
            def start(self, group: Group):
                gevent.sleep(0.5)
                return super().start(group)

            @task
            def my_task(self):
                gevent.sleep(600)

Review comment (bug):
test would fail occasionally with state being `stopping` instead of `stopped`.

Code after:
                self.assertEqual(
                    36, test_shape.get_current_user_count(), "Shape is not seeing stage 2 runner user count correctly"
                )
            self.assertDictEqual(
                master.reported_user_classes_count, {"FixedUser1": 1, "FixedUser2": 11, "TestUser": 24}
            )

            # Ensure new stage with less users has been reached
            # and expected count of the fixed users is present
            sleep(1)
            for worker in workers:
                self.assertEqual(4, worker.user_count, "Shape test has not reached stage 3")
                self.assertEqual(
                    12, test_shape.get_current_user_count(), "Shape is not seeing stage 3 runner user count correctly"
                )
            self.assertDictEqual(master.reported_user_classes_count, {"FixedUser1": 1, "FixedUser2": 11, "TestUser": 0})

            # Ensure test stops at the end
            sleep(0.5)
            for worker in workers:
                self.assertEqual(0, worker.user_count, "Shape test has not stopped")
                self.assertEqual(
                    0, test_shape.get_current_user_count(), "Shape is not seeing stopped runner user count correctly"
                )
            self.assertDictEqual(master.reported_user_classes_count, {"FixedUser1": 0, "FixedUser2": 0, "TestUser": 0})
            sleep(1.0)
            self.assertEqual(STATE_STOPPED, master.state)

    def test_distributed_shape_with_stop_timeout(self):
        """
        Full integration test that starts both a MasterRunner and five WorkerRunner instances
        and tests a basic LoadTestShape with scaling up and down users
        """

        class TestUser1(User):
            def start(self, group: Group):
                gevent.sleep(0.5)
                return super().start(group)

            @task
            def my_task(self):
                gevent.sleep(0)

        class TestUser2(User):
            def start(self, group: Group):
                gevent.sleep(0.5)
                return super().start(group)

            @task
            def my_task(self):
                gevent.sleep(600)