Merge pull request #1327 from locustio/environment-refactor

Add Runners, WebUI and Environment to the public API

docs/api.rst

@@ -105,3 +105,23 @@ The event hooks are instances of the **locust.events.EventHook** class:
 
     It's highly recommended that you add a wildcard keyword argument in your event listeners
     to prevent your code from breaking if new arguments are added in a future version.
+
+
+Locust Runner classes
+=====================
+
+.. autoclass:: locust.runners.LocustRunner
+    :members: start, stop, quit, user_count
+
+.. autoclass:: locust.runners.LocalLocustRunner
+
+.. autoclass:: locust.runners.MasterLocustRunner
+
+.. autoclass:: locust.runners.WorkerLocustRunner
+
+
+Web UI class
+============
+
+.. autoclass:: locust.web.WebUI
+    :members:

docs/use-as-lib.rst

@@ -2,45 +2,41 @@
 Using Locust as a library
 ==========================
 
-It's possible to use Locust as a library instead of running Locust by invoking the ``locust`` command.
+It's possible to use Locust as a library, instead of running Locust using the ``locust`` command.
 
-Here's an example::
+To run Locust as a library you need to create an :py:class:`Environment <locust.env.Environment>` instance:
+
+.. code-block:: python
 
-    import gevent
-    from locust import HttpLocust, TaskSet, task, between
-    from locust.runners import LocalLocustRunner
     from locust.env import Environment
-    from locust.stats import stats_printer
-    from locust.log import setup_logging
-    from locust.web import WebUI
-    
-    setup_logging("INFO", None)
-    
-    
-    class User(HttpLocust):
-        wait_time = between(1, 3)
-        host = "https://docs.locust.io"
-        
-        class task_set(TaskSet):
-            @task
-            def my_task(self):
-                self.client.get("/")
-            
-            @task
-            def task_404(self):
-                self.client.get("/non-existing-path")
     
-    # setup Environment and Runner
-    env = Environment()
-    runner = LocalLocustRunner(environment=env, locust_classes=[User])
-    # start a WebUI instance
-    web_ui = WebUI(environment=env)
-    gevent.spawn(lambda: web_ui.start("127.0.0.1", 8089))
-    
-    # start a greenlet that periodically outputs the current stats
-    gevent.spawn(stats_printer(env.stats))
-    
-    # start the test
-    runner.start(1, hatch_rate=10)
-    # wait for the greenlets (indefinitely)
-    runner.greenlet.join()
+    env = Environment(locust_classes=[MyTestUser])
+
+The :py:class:`Environment <locust.env.Environment>` instance's 
+:py:meth:`create_local_runner <locust.env.Environment.create_local_runner>`, 
+:py:meth:`create_master_runner <locust.env.Environment.create_master_runner>` or 
+:py:meth:`create_worker_runner <locust.env.Environment.create_worker_runner> can then be used to start a 
+:py:class:`LocustRunner <locust.runners.LocustRunner>` instance, which can be used to start a load test:
+
+.. code-block:: python
+
+    env.create_local_runner()
+    env.runner.start(5000, hatch_rate=20)
+    env.runner.greenlet.join()
+
+We could also use the :py:class:`Environment <locust.env.Environment>` instance's 
+:py:meth:`create_web_ui <locust.env.Environment.create_web_ui>` method to start a Web UI that can be used 
+to view the stats, and to control the runner (e.g. start and stop load tests):
+
+.. code-block:: python
+
+    env.create_local_runner()
+    env.create_web_ui()
+    env.web_ui.greenlet.join()
+
+
+Full example
+============
+
+.. literalinclude:: ../examples/use_as_lib.py
+    :language: python

examples/use_as_lib.py

@@ -1,44 +1,42 @@
 import gevent
-from locust import HttpLocust, TaskSet, task, between
-from locust.runners import LocalLocustRunner
+from locust import HttpLocust, task, between
 from locust.env import Environment
 from locust.stats import stats_printer
 from locust.log import setup_logging
-from locust.web import WebUI
 
 setup_logging("INFO", None)
 
 
 class User(HttpLocust):
     wait_time = between(1, 3)
     host = "https://docs.locust.io"
-
-    class task_set(TaskSet):
-        @task
-        def my_task(self):
-            self.client.get("/")
-
-        @task
-        def task_404(self):
-            self.client.get("/non-existing-path")
 
-# setup Environment and Runner
-env = Environment()
-runner = LocalLocustRunner(environment=env, locust_classes=[User])
-# start a WebUI instance
-web_ui = WebUI(environment=env)
-gevent.spawn(lambda: web_ui.start("127.0.0.1", 8089))
+    @task
+    def my_task(self):
+        self.client.get("/")
 
+    @task
+    def task_404(self):
+        self.client.get("/non-existing-path")
+
+# setup Environment and Runner
+env = Environment(locust_classes=[User])
+env.create_local_runner()
 
-# TODO: fix 
-#def on_request_success(request_type, name, response_time, response_length, **kwargs):
-#    report_to_grafana("%_%s" % (request_type, name), response_time)
-#env.events.request_succes.add_listener(on_request_success)
+# start a WebUI instance
+env.create_web_ui("127.0.0.1", 8089)
 
 # start a greenlet that periodically outputs the current stats
-gevent.spawn(stats_printer(runner.stats))
+gevent.spawn(stats_printer(env.stats))
 
 # start the test
-runner.start(1, hatch_rate=10)
-# wait for the greenlets (indefinitely)
-runner.greenlet.join()
+env.runner.start(1, hatch_rate=10)
+
+# in 60 seconds stop the runner
+gevent.spawn_later(60, lambda: env.runner.quit())
+
+# wait for the greenlets
+env.runner.greenlet.join()
+
+# stop the web server for good measures
+env.web_ui.stop()

locust/env.py

@@ -1,4 +1,8 @@
 from .event import Events
+from .exception import RunnerAlreadyExistsError
+from .stats import RequestStats
+from .runners import LocalLocustRunner, MasterLocustRunner, WorkerLocustRunner
+from .web import WebUI
 
 
 class Environment:
@@ -8,15 +12,18 @@ class Environment:
     See :ref:`events` for available events.
     """
 
+    locust_classes = []
+    """Locust User classes that the runner will run"""
+
+    stats = None
+    """Reference to RequestStats instance"""
+
     runner = None
-    """Reference to the LocustRunner instance"""
+    """Reference to the :class:`LocustRunner <locust.runners.LocustRunner>` instance"""
 
     web_ui = None
     """Reference to the WebUI instance"""
 
-    options = None
-    """Parsed command line options"""
-
     host = None
     """Base URL of the target system"""
 
@@ -38,22 +45,10 @@ class Environment:
     If False, exeptions will be raised.
     """
 
-    master_host = "127.0.0.1"
-    """Hostname of master node that the worker should connect to"""
-
-    master_port = 5557
-    """Port of master node that the worker should connect to. Defaults to 5557."""
-
-    master_bind_host = "*"
-    """Hostname/interfaces that the master node should expect workers to connect to. Defaults to '*' which means all interfaces."""
-
-    master_bind_port = 5557    
-    """Port that the master node should listen to and expect workers to connect to. Defaults to 5557."""
-
     def  __init__(
-        self, 
+        self, *,
+        locust_classes=[],
         events=None, 
-        options=None, 
         host=None, 
         reset_stats=False, 
         step_load=False, 
@@ -65,10 +60,64 @@ def  __init__(
         else:
             self.events = Events()
 
-        self.options = options
+        self.locust_classes = locust_classes
+        self.stats = RequestStats()
         self.host = host
         self.reset_stats = reset_stats
         self.step_load = step_load
         self.stop_timeout = stop_timeout
         self.catch_exceptions = catch_exceptions
+
+    def _create_runner(self, runner_class, *args, **kwargs):
+        if self.runner is not None:
+            raise RunnerAlreadyExistsError("Environment.runner already exists (%s)" % self.runner)
+        self.runner = runner_class(self, *args, **kwargs)
+        return self.runner
+
+    def create_local_runner(self):
+        """
+        Create a :class:`LocalLocustRunner <locust.runners.LocalLocustRunner>` instance for this Environment
+        """
+        return self._create_runner(LocalLocustRunner)
+
+    def create_master_runner(self, master_bind_host="*", master_bind_port=5557):
+        """
+        Create a :class:`MasterLocustRunner <locust.runners.MasterLocustRunner>` instance for this Environment
+        
+        :param master_bind_host: Interface/host that the master should use for incoming worker connections. 
+                                 Defaults to "*" which means all interfaces.
+        :param master_bind_port: Port that the master should listen for incoming worker connections on
+        """
+        return self._create_runner(
+            MasterLocustRunner,
+            master_bind_host=master_bind_host,
+            master_bind_port=master_bind_port,
+        )
+
+    def create_worker_runner(self, master_host, master_port):
+        """
+        Create a :class:`WorkerLocustRunner <locust.runners.WorkerLocustRunner>` instance for this Environment
+        
+        :param master_host: Host/IP of a running master node
+        :param master_port: Port on master node to connect to
+        """
+        # Create a new RequestStats with use_response_times_cache set to False to save some memory
+        # and CPU cycles, since the response_times_cache is not needed for Worker nodes
+        self.stats = RequestStats(use_response_times_cache=False)
+        return self._create_runner(
+            WorkerLocustRunner,
+            master_host=master_host,
+            master_port=master_port,
+        )
+
+    def create_web_ui(self, host="*", port=8089, auth_credentials=None):
+        """
+        Creates a :class:`WebUI <locust.web.WebUI>` instance for this Environment and start running the web server
         
+        :param host: Host/interface that the web server should accept connections to. Defaults to "*"
+                     which means all interfaces
+        :param port: Port that the web server should listen to
+        :param auth_credentials: If provided (in format "username:password") basic auth will be enabled
+        """
+        self.web_ui = WebUI(self, host, port, auth_credentials=auth_credentials)
+        return self.web_ui

locust/exception.py

@@ -50,4 +50,7 @@ class AuthCredentialsError(ValueError):
     Exception when the auth credentials provided
     are not in the correct format
     """
-    pass
+    pass
+
+class RunnerAlreadyExistsError(Exception):
+    pass

locust/main.py

@@ -16,11 +16,9 @@
 from .env import Environment
 from .inspectlocust import get_task_ratio_dict, print_task_ratio
 from .log import setup_logging, greenlet_exception_logger
-from .runners import LocalLocustRunner, MasterLocustRunner, WorkerLocustRunner
 from .stats import (print_error_report, print_percentile_stats, print_stats,
                     stats_printer, stats_writer, write_csv_files)
 from .util.timespan import parse_timespan
-from .web import WebUI
 from .exception import AuthCredentialsError
 
 _internals = [Locust, HttpLocust]
@@ -89,14 +87,14 @@ def __import_locustfile__(filename, path):
     return imported.__doc__, locusts
 
 
-def create_environment(options, events=None):
+def create_environment(locust_classes, options, events=None):
     """
     Create an Environment instance from options
     """
     return Environment(
+        locust_classes=locust_classes,
         events=events,
         host=options.host,
-        options=options,
         reset_stats=options.reset_stats,
         step_load=options.step_load,
         stop_timeout=options.stop_timeout,
@@ -153,7 +151,7 @@ def main():
         locust_classes = list(locusts.values())
 
     # create locust Environment
-    environment = create_environment(options, events=locust.events)
+    environment = create_environment(locust_classes, options, events=locust.events)
 
     if options.show_task_ratio:
         print("\n Task ratio per locust class")
@@ -186,25 +184,18 @@ def main():
             sys.exit(1)
 
     if options.master:
-        runner = MasterLocustRunner(
-            environment, 
-            locust_classes,
-            master_bind_host=options.master_bind_host,
+        runner = environment.create_master_runner(
+            master_bind_host=options.master_bind_host, 
             master_bind_port=options.master_bind_port,
         )
     elif options.worker:
         try:
-            runner = WorkerLocustRunner(
-                environment, 
-                locust_classes,
-                master_host=options.master_host,
-                master_port=options.master_port,
-            )
+            runner = environment.create_worker_runner(options.master_host, options.master_port)
         except socket.error as e:
             logger.error("Failed to connect to the Locust master: %s", e)
             sys.exit(-1)
     else:
-        runner = LocalLocustRunner(environment, locust_classes)
+        runner = environment.create_local_runner()
 
     # main_greenlet is pointing to runners.greenlet by default, it will point the web greenlet later if in web mode
     main_greenlet = runner.greenlet
@@ -233,7 +224,7 @@ def timelimit_stop():
         # spawn web greenlet
         logger.info("Starting web monitor at http://%s:%s" % (options.web_host or "*", options.web_port))
         try:
-            web_ui = WebUI(environment=environment, auth_credentials=options.web_auth)
+            web_ui = environment.create_web_ui(auth_credentials=options.web_auth)
         except AuthCredentialsError:
             logger.error("Credentials supplied with --web-auth should have the format: username:password")
             sys.exit(1)