Make Runner.spawn_users() return the new User instances for use with …

…locust-as-a-library (PR #1791)

docs/use-as-lib.rst

@@ -26,6 +26,17 @@ The :py:class:`Environment <locust.env.Environment>` instance's
     env.runner.start(5000, spawn_rate=20)
     env.runner.greenlet.join()
 
+It is also possible to bypass the dispatch and distribution logic, and manually control the spawned users:
+
+.. code-block:: python
+
+    new_users = env.runner.spawn_users({MyUserClass.__name__: 2})
+    new_users[1].my_custom_token = "custom-token-2"
+    new_users[0].my_custom_token = "custom-token-1"
+
+The above example only works on standalone mode and is an experimental feature, meaning that it could be
+removed in future versions. But it's useful if you want to have fine-grained control on the spawned users. 
+
 .. note::
 
     Do not attempt to create a master runner and worker(s) in the same Python process. It does not work, and even if it did, it would not give you better performance than running a single LocalRunner. Every worker *must* run in its own process, there is no way around that.

locust/runners.py

@@ -205,20 +205,25 @@ def spawn_users(self, user_classes_spawn_count: Dict[str, int], wait: bool = Fal
 
         def spawn(user_class: str, spawn_count: int):
             n = 0
+            new_users = []
             while n < spawn_count:
                 new_user = self.user_classes_by_name[user_class](self.environment)
                 new_user.start(self.user_greenlets)
+                new_users.append(new_user)
                 n += 1
                 if n % 10 == 0 or n == spawn_count:
                     logger.debug("%i users spawned" % self.user_count)
             logger.debug("All users of class %s spawned" % user_class)
+            return new_users
 
+        new_users = []
         for user_class, spawn_count in user_classes_spawn_count.items():
-            spawn(user_class, spawn_count)
+            new_users += spawn(user_class, spawn_count)
 
         if wait:
             self.user_greenlets.join()
             logger.info("All users stopped\n")
+        return new_users
 
     def stop_users(self, user_classes_stop_count: Dict[str, int]):
         async_calls_to_stop = Group()

locust/test/test_runners.py

@@ -151,7 +151,8 @@ def trigger(self):
                     triggered[0] = True
 
         runner = Environment(user_classes=[BaseUser]).create_local_runner()
-        runner.spawn_users({BaseUser.__name__: 2}, wait=False)
+        users = runner.spawn_users({BaseUser.__name__: 2}, wait=False)
+        self.assertEqual(2, len(users))
         self.assertEqual(2, len(runner.user_greenlets))
         g1 = list(runner.user_greenlets)[0]
         g2 = list(runner.user_greenlets)[1]
@@ -339,7 +340,8 @@ def on_stop(self):
                 BaseUser.stop_triggered = True
 
         runner = Environment(user_classes=[BaseUser]).create_local_runner()
-        runner.spawn_users({BaseUser.__name__: 1}, wait=False)
+        users = runner.spawn_users({BaseUser.__name__: 1}, wait=False)
+        self.assertEqual(1, len(users))
         timeout = gevent.Timeout(0.5)
         timeout.start()
         try:
@@ -364,7 +366,8 @@ def on_stop(self):
                 BaseUser.stop_count += 1
 
         runner = Environment(user_classes=[BaseUser]).create_local_runner()
-        runner.spawn_users({BaseUser.__name__: 10}, wait=False)
+        users = runner.spawn_users({BaseUser.__name__: 10}, wait=False)
+        self.assertEqual(10, len(users))
         timeout = gevent.Timeout(0.3)
         timeout.start()
         try: