Introduce LOCUST_SPAWN_AT_LEAST_ONE_USER_OF_EACH_USER_CLASS

locust/distribution.py

@@ -1,4 +1,5 @@
 import math
+import os
 from itertools import combinations_with_replacement
 from operator import attrgetter
 from typing import Dict, List, Type
@@ -23,9 +24,16 @@ def weight_users(user_classes: List[Type[User]], user_count: int) -> Dict[str, i
 
     user_classes_count = {user_class.__name__: 0 for user_class in user_classes}
 
-    # If the number of users is less than the number of user classes, at most one user of each user class
-    # is chosen. User classes with higher weight are chosen first.
-    if user_count <= len(user_classes):
+    if user_count == 0:
+        return user_classes_count
+
+    spawn_at_least_one_user_of_each_user_class = {"true": True, "false": False}[
+        os.getenv("LOCUST_SPAWN_AT_LEAST_ONE_USER_OF_EACH_USER_CLASS", "true").lower()
+    ]
+
+    if spawn_at_least_one_user_of_each_user_class and user_count <= len(user_classes):
+        # If the number of users is less than the number of user classes, at most one user of each user class
+        # is chosen. User classes with higher weight are chosen first.
         user_classes_count.update(
             {
                 user_class.__name__: 1
@@ -34,25 +42,37 @@ def weight_users(user_classes: List[Type[User]], user_count: int) -> Dict[str, i
         )
         return user_classes_count
 
-    # If the number of users is greater than or equal to the number of user classes, at least one user of each
-    # user class will be chosen. The greater number of users is, the better the actual distribution
-    # of users will match the desired one (as dictated by the weight attributes).
-    weights = list(map(attrgetter("weight"), user_classes))
-    relative_weights = [weight / sum(weights) for weight in weights]
-    user_classes_count = {
-        user_class.__name__: round(relative_weight * user_count) or 1
-        for user_class, relative_weight in zip(user_classes, relative_weights)
-    }
+    user_classes_count = _approximate_user_classes_count(
+        user_classes, user_count, spawn_at_least_one_user_of_each_user_class
+    )
 
     if sum(user_classes_count.values()) == user_count:
         return user_classes_count
 
+    user_classes_count = _find_ideal_users_to_add_or_remove(
+        user_classes, user_count - sum(user_classes_count.values()), user_classes_count
+    )
+    assert sum(user_classes_count.values()) == user_count
+    return user_classes_count
+
+
+def _approximate_user_classes_count(
+    user_classes: List[Type[User]], user_count: int, spawn_at_least_one_user_of_each_user_class: bool
+) -> Dict[str, int]:
+    if spawn_at_least_one_user_of_each_user_class:
+        # If the number of users is greater than or equal to the number of user classes, at least one user of each
+        # user class will be chosen. The greater number of users is, the better the actual distribution
+        # of users will match the desired one (as dictated by the weight attributes).
+        min_user_class_count = 1
     else:
-        user_classes_count = _find_ideal_users_to_add_or_remove(
-            user_classes, user_count - sum(user_classes_count.values()), user_classes_count
-        )
-        assert sum(user_classes_count.values()) == user_count
-        return user_classes_count
+        min_user_class_count = 0
+
+    weights = list(map(attrgetter("weight"), user_classes))
+    relative_weights = [weight / sum(weights) for weight in weights]
+    return {
+        user_class.__name__: round(relative_weight * user_count) or min_user_class_count
+        for user_class, relative_weight in zip(user_classes, relative_weights)
+    }
 
 
 def _find_ideal_users_to_add_or_remove(

locust/test/test_distribution.py

@@ -3,9 +3,10 @@
 
 from locust import User
 from locust.distribution import weight_users
+from locust.test.util import patch_env
 
 
-class TestDistribution(unittest.TestCase):
+class TestDistributionSpawnAtLeastOneUserOfEachUserClass(unittest.TestCase):
     def test_distribution_no_user_classes(self):
         user_classes_count = weight_users(user_classes=[], user_count=0)
         self.assertDictEqual(user_classes_count, {})
@@ -101,6 +102,238 @@ class User3(User):
         user_classes_count = weight_users(user_classes=[User1, User2, User3], user_count=11)
         self.assertDictEqual(user_classes_count, {"User1": 2, "User2": 4, "User3": 5})
 
+        user_classes_count = weight_users(user_classes=[User1, User2, User3], user_count=30)
+        self.assertDictEqual(user_classes_count, {"User1": 5, "User2": 10, "User3": 15})
+
+        user_classes_count = weight_users(user_classes=[User1, User2, User3], user_count=300)
+        self.assertDictEqual(user_classes_count, {"User1": 50, "User2": 100, "User3": 150})
+
+    def test_distribution_unequal_and_non_unique_weights_and_fewer_amount_than_user_classes(self):
+        class User1(User):
+            weight = 1
+
+        class User2(User):
+            weight = 2
+
+        class User3(User):
+            weight = 2
+
+        user_classes_count = weight_users(user_classes=[User1, User2, User3], user_count=0)
+        self.assertDictEqual(user_classes_count, {"User1": 0, "User2": 0, "User3": 0})
+
+        user_classes_count = weight_users(user_classes=[User1, User2, User3], user_count=1)
+        self.assertDictEqual(user_classes_count, {"User1": 0, "User2": 1, "User3": 0})
+
+        user_classes_count = weight_users(user_classes=[User1, User2, User3], user_count=2)
+        self.assertDictEqual(user_classes_count, {"User1": 0, "User2": 1, "User3": 1})
+
+    def test_distribution_unequal_and_non_unique_weights(self):
+        class User1(User):
+            weight = 1
+
+        class User2(User):
+            weight = 2
+
+        class User3(User):
+            weight = 2
+
+        user_classes_count = weight_users(user_classes=[User1, User2, User3], user_count=3)
+        self.assertDictEqual(user_classes_count, {"User1": 1, "User2": 1, "User3": 1})
+
+        user_classes_count = weight_users(user_classes=[User1, User2, User3], user_count=4)
+        self.assertDictEqual(user_classes_count, {"User1": 1, "User2": 1, "User3": 2})
+
+        user_classes_count = weight_users(user_classes=[User1, User2, User3], user_count=5)
+        self.assertDictEqual(user_classes_count, {"User1": 1, "User2": 2, "User3": 2})
+
+        user_classes_count = weight_users(user_classes=[User1, User2, User3], user_count=6)
+        self.assertDictEqual(user_classes_count, {"User1": 1, "User2": 3, "User3": 2})
+
+        user_classes_count = weight_users(user_classes=[User1, User2, User3], user_count=10)
+        self.assertDictEqual(user_classes_count, {"User1": 2, "User2": 4, "User3": 4})
+
+        user_classes_count = weight_users(user_classes=[User1, User2, User3], user_count=11)
+        self.assertDictEqual(user_classes_count, {"User1": 2, "User2": 5, "User3": 4})
+
+    def test_distribution_large_number_of_users(self):
+        class User1(User):
+            weight = 5
+
+        class User2(User):
+            weight = 55
+
+        class User3(User):
+            weight = 37
+
+        class User4(User):
+            weight = 2
+
+        class User5(User):
+            weight = 97
+
+        class User6(User):
+            weight = 41
+
+        class User7(User):
+            weight = 33
+
+        class User8(User):
+            weight = 19
+
+        class User9(User):
+            weight = 19
+
+        class User10(User):
+            weight = 34
+
+        class User11(User):
+            weight = 78
+
+        class User12(User):
+            weight = 76
+
+        class User13(User):
+            weight = 28
+
+        class User14(User):
+            weight = 62
+
+        class User15(User):
+            weight = 69
+
+        for user_count in range(1044523783783, 1044523783783 + 1000):
+            ts = time.perf_counter()
+            user_classes_count = weight_users(
+                user_classes=[
+                    User1,
+                    User2,
+                    User3,
+                    User4,
+                    User5,
+                    User6,
+                    User7,
+                    User8,
+                    User9,
+                    User10,
+                    User11,
+                    User12,
+                    User13,
+                    User14,
+                    User15,
+                ],
+                user_count=user_count,
+            )
+            delta_ms = 1e3 * (time.perf_counter() - ts)
+            self.assertEqual(sum(user_classes_count.values()), user_count)
+            self.assertLessEqual(delta_ms, 100)
+
+
+class TestDistributionDoNotSpawnAtLeastOneUserOfEachUserClass(unittest.TestCase):
+    def run(self, result=None):
+        with patch_env("LOCUST_SPAWN_AT_LEAST_ONE_USER_OF_EACH_USER_CLASS", "false"):
+            super(TestDistributionDoNotSpawnAtLeastOneUserOfEachUserClass, self).run(result)
+
+    def test_distribution_no_user_classes(self):
+        user_classes_count = weight_users(user_classes=[], user_count=0)
+        self.assertDictEqual(user_classes_count, {})
+
+        user_classes_count = weight_users(user_classes=[], user_count=1)
+        self.assertDictEqual(user_classes_count, {})
+
+    def test_distribution_equal_weights_and_fewer_amount_than_user_classes(self):
+        class User1(User):
+            weight = 1
+
+        class User2(User):
+            weight = 1
+
+        class User3(User):
+            weight = 1
+
+        user_classes_count = weight_users(user_classes=[User1, User2, User3], user_count=0)
+        self.assertDictEqual(user_classes_count, {"User1": 0, "User2": 0, "User3": 0})
+
+        user_classes_count = weight_users(user_classes=[User1, User2, User3], user_count=1)
+        self.assertDictEqual(user_classes_count, {"User1": 1, "User2": 0, "User3": 0})
+
+        user_classes_count = weight_users(user_classes=[User1, User2, User3], user_count=2)
+        self.assertDictEqual(user_classes_count, {"User1": 0, "User2": 1, "User3": 1})
+
+    def test_distribution_equal_weights(self):
+        class User1(User):
+            weight = 1
+
+        class User2(User):
+            weight = 1
+
+        class User3(User):
+            weight = 1
+
+        user_classes_count = weight_users(user_classes=[User1, User2, User3], user_count=3)
+        self.assertDictEqual(user_classes_count, {"User1": 1, "User2": 1, "User3": 1})
+
+        user_classes_count = weight_users(user_classes=[User1, User2, User3], user_count=4)
+        self.assertDictEqual(user_classes_count, {"User1": 2, "User2": 1, "User3": 1})
+
+        user_classes_count = weight_users(user_classes=[User1, User2, User3], user_count=5)
+        self.assertDictEqual(user_classes_count, {"User1": 1, "User2": 2, "User3": 2})
+
+        user_classes_count = weight_users(user_classes=[User1, User2, User3], user_count=6)
+        self.assertDictEqual(user_classes_count, {"User1": 2, "User2": 2, "User3": 2})
+
+    def test_distribution_unequal_and_unique_weights_and_fewer_amount_than_user_classes(self):
+        class User1(User):
+            weight = 1
+
+        class User2(User):
+            weight = 2
+
+        class User3(User):
+            weight = 3
+
+        user_classes_count = weight_users(user_classes=[User1, User2, User3], user_count=0)
+        self.assertDictEqual(user_classes_count, {"User1": 0, "User2": 0, "User3": 0})
+
+        user_classes_count = weight_users(user_classes=[User1, User2, User3], user_count=1)
+        self.assertDictEqual(user_classes_count, {"User1": 0, "User2": 0, "User3": 1})
+
+        user_classes_count = weight_users(user_classes=[User1, User2, User3], user_count=2)
+        self.assertDictEqual(user_classes_count, {"User1": 0, "User2": 1, "User3": 1})
+
+    def test_distribution_unequal_and_unique_weights(self):
+        class User1(User):
+            weight = 1
+
+        class User2(User):
+            weight = 2
+
+        class User3(User):
+            weight = 3
+
+        user_classes_count = weight_users(user_classes=[User1, User2, User3], user_count=3)
+        self.assertDictEqual(user_classes_count, {"User1": 0, "User2": 1, "User3": 2})
+
+        user_classes_count = weight_users(user_classes=[User1, User2, User3], user_count=4)
+        self.assertDictEqual(user_classes_count, {"User1": 1, "User2": 1, "User3": 2})
+
+        user_classes_count = weight_users(user_classes=[User1, User2, User3], user_count=5)
+        self.assertDictEqual(user_classes_count, {"User1": 1, "User2": 2, "User3": 2})
+
+        user_classes_count = weight_users(user_classes=[User1, User2, User3], user_count=6)
+        self.assertDictEqual(user_classes_count, {"User1": 1, "User2": 2, "User3": 3})
+
+        user_classes_count = weight_users(user_classes=[User1, User2, User3], user_count=10)
+        self.assertDictEqual(user_classes_count, {"User1": 2, "User2": 3, "User3": 5})
+
+        user_classes_count = weight_users(user_classes=[User1, User2, User3], user_count=11)
+        self.assertDictEqual(user_classes_count, {"User1": 2, "User2": 4, "User3": 5})
+
+        user_classes_count = weight_users(user_classes=[User1, User2, User3], user_count=30)
+        self.assertDictEqual(user_classes_count, {"User1": 5, "User2": 10, "User3": 15})
+
+        user_classes_count = weight_users(user_classes=[User1, User2, User3], user_count=300)
+        self.assertDictEqual(user_classes_count, {"User1": 50, "User2": 100, "User3": 150})
+
     def test_distribution_unequal_and_non_unique_weights_and_fewer_amount_than_user_classes(self):
         class User1(User):
             weight = 1

locust/test/test_runners.py

@@ -46,6 +46,7 @@
     User,
     task,
 )
+from .util import patch_env
 
 NETWORK_BROKEN = "network broken"
 
@@ -843,7 +844,7 @@ def tick(self):
                     return None
 
         locust_worker_additional_wait_before_ready_after_stop = 5
-        with mock.patch("locust.runners.WORKER_REPORT_INTERVAL", new=0.3), _patch_env(
+        with mock.patch("locust.runners.WORKER_REPORT_INTERVAL", new=0.3), patch_env(
             "LOCUST_WORKER_ADDITIONAL_WAIT_BEFORE_READY_AFTER_STOP",
             str(locust_worker_additional_wait_before_ready_after_stop),
         ):
@@ -1148,7 +1149,7 @@ def tick(self):
             user_class.weight = random.uniform(1, 20)
 
         locust_worker_additional_wait_before_ready_after_stop = 5
-        with mock.patch("locust.runners.WORKER_REPORT_INTERVAL", new=0.3), _patch_env(
+        with mock.patch("locust.runners.WORKER_REPORT_INTERVAL", new=0.3), patch_env(
             "LOCUST_WORKER_ADDITIONAL_WAIT_BEFORE_READY_AFTER_STOP",
             str(locust_worker_additional_wait_before_ready_after_stop),
         ):
@@ -1266,7 +1267,7 @@ def tick(self):
                     return None
 
         locust_worker_additional_wait_before_ready_after_stop = 2
-        with mock.patch("locust.runners.WORKER_REPORT_INTERVAL", new=0.3), _patch_env(
+        with mock.patch("locust.runners.WORKER_REPORT_INTERVAL", new=0.3), patch_env(
             "LOCUST_WORKER_ADDITIONAL_WAIT_BEFORE_READY_AFTER_STOP",
             str(locust_worker_additional_wait_before_ready_after_stop),
         ):
@@ -2232,12 +2233,12 @@ def assert_cache_hits():
         assert_cache_hits()
 
         master._wait_for_workers_report_after_ramp_up.cache_clear()
-        with _patch_env("LOCUST_WAIT_FOR_WORKERS_REPORT_AFTER_RAMP_UP", "5.7"):
+        with patch_env("LOCUST_WAIT_FOR_WORKERS_REPORT_AFTER_RAMP_UP", "5.7"):
             self.assertEqual(master._wait_for_workers_report_after_ramp_up(), 5.7)
             assert_cache_hits()
 
         master._wait_for_workers_report_after_ramp_up.cache_clear()
-        with mock.patch("locust.runners.WORKER_REPORT_INTERVAL", new=1.5), _patch_env(
+        with mock.patch("locust.runners.WORKER_REPORT_INTERVAL", new=1.5), patch_env(
             "LOCUST_WAIT_FOR_WORKERS_REPORT_AFTER_RAMP_UP", "5.7 * WORKER_REPORT_INTERVAL"
         ):
             self.assertEqual(master._wait_for_workers_report_after_ramp_up(), 5.7 * 1.5)
@@ -2246,19 +2247,6 @@ def assert_cache_hits():
         master._wait_for_workers_report_after_ramp_up.cache_clear()
 
 
-@contextmanager
-def _patch_env(name: str, value: str):
-    prev_value = os.getenv(name)
-    os.environ[name] = value
-    try:
-        yield
-    finally:
-        if prev_value is None:
-            del os.environ[name]
-        else:
-            os.environ[name] = prev_value
-
-
 class TestWorkerRunner(LocustTestCase):
     def setUp(self):
         super().setUp()