Merge pull request #1167 from cyberw/get-parameters-from-conf-file-an…

…d-env-vars

Use ConfigArgParse instead of argparse, to support getting parameters from config file and/or env vars.

docs/quickstart.rst

@@ -147,6 +147,22 @@ host defaults to 127.0.0.1):
     $ locust -f locust_files/my_locust_file.py --slave --master-host=192.168.0.100
 
 
+Parameters can also be set in a `config file <https://github.com/bw2/ConfigArgParse#config-file-syntax>`_ (locust.conf or ~/.locust.conf) or in env vars, prefixed by LOCUST\_
+
+For example: (this will do the same thing as the previous command)
+
+.. code-block::
+
+    # locust.conf in current directory
+    locustfile locust_files/my_locust_file.py
+    slave
+
+
+.. code-block:: console
+
+    $ LOCUST_MASTER_HOST=192.168.0.100 locust
+
+
 .. note::
 
     To see all available options type: ``locust --help``

locust/main.py

@@ -6,7 +6,7 @@
 import socket
 import sys
 import time
-import argparse
+import configargparse
 
 import gevent
 
@@ -25,15 +25,14 @@
 version = locust.__version__
 
 
-def parse_options(args=None):
+def parse_options(args=None, default_config_files=['~/.locust.conf','locust.conf']):
     """
-    Handle command-line options with argparse.ArgumentParser.
+    Handle command-line options with configargparse.ArgumentParser.
 
-    Return list of arguments, largely for use in `parse_arguments`.
+    Returns a two-tuple of parser + the output from parse_args()
     """
-
     # Initialize
-    parser = argparse.ArgumentParser()
+    parser = configargparse.ArgumentParser(default_config_files=default_config_files, auto_env_var_prefix="LOCUST_", add_env_var_help=False)
 
     parser.add_argument(
         '-H', '--host',
@@ -132,7 +131,7 @@ def parse_options(args=None):
     parser.add_argument(
         '--no-web',
         action='store_true',
-        help="Disable the web interface, and instead start running the test immediately. Requires -c and -r to be specified."
+        help="Disable the web interface, and instead start running the test immediately. Requires -c and -t to be specified."
     )
 
     # Number of clients
@@ -258,8 +257,6 @@ def parse_options(args=None):
         metavar='LocustClass',
     )
 
-    # Finalize
-    # Return two-tuple of parser + the output from parse_args
     return parser, parser.parse_args(args=args)
 
 
@@ -458,7 +455,7 @@ def timelimit_stop():
 
     if not options.no_web and not options.slave:
         # spawn web greenlet
-        logger.info("Starting web monitor at %s:%s" % (options.web_host or "*", options.port))
+        logger.info("Starting web monitor at http://%s:%s" % (options.web_host or "*", options.port))
         main_greenlet = gevent.spawn(web.start, locust_classes, options)
 
     if not options.master and not options.slave:

locust/test/test_parser.py

@@ -1,11 +1,13 @@
 import unittest
+import os
+import tempfile
 
 from locust.main import parse_options
 
 
 class TestParser(unittest.TestCase):
     def setUp(self):
-        self.parser, _ = parse_options()
+        self.parser, _ = parse_options(default_config_files=[])
 
     def test_default(self):
         opts = self.parser.parse_args([])
@@ -32,3 +34,15 @@ def test_skip_log_setup(self):
         ]
         opts = self.parser.parse_args(args)
         self.assertEqual(opts.skip_log_setup, True)
+
+    def test_parameter_parsing(self):
+        with tempfile.NamedTemporaryFile(mode='w') as file:
+            os.environ['LOCUST_LOCUSTFILE'] = "locustfile_from_env"
+            file.write("host host_from_config\nweb-host webhost_from_config")
+            file.flush()
+            parser, _ = parse_options(default_config_files=[file.name])
+            options = parser.parse_args(['-H','host_from_args'])
+            del os.environ['LOCUST_LOCUSTFILE']
+        self.assertEqual(options.web_host, 'webhost_from_config')
+        self.assertEqual(options.locustfile, 'locustfile_from_env')
+        self.assertEqual(options.host, 'host_from_args') # overridden

locust/test/test_web.py

@@ -26,7 +26,7 @@ def setUp(self):
         super(TestWebUI, self).setUp()
 
         stats.global_stats.clear_all()
-        parser = parse_options()[0]
+        parser = parse_options(default_config_files=[])[0]
         self.options = parser.parse_args([])
         runners.locust_runner = LocustRunner([], self.options)
 

locust/test/test_zmqrpc.py

@@ -8,7 +8,7 @@
 class ZMQRPC_tests(LocustTestCase):
     def setUp(self):
         super(ZMQRPC_tests, self).setUp()
-        self.server = zmqrpc.Server('*', 0)
+        self.server = zmqrpc.Server('127.0.0.1', 0)
         self.client = zmqrpc.Client('localhost', self.server.port, 'identity')
 
     def tearDown(self):
@@ -40,7 +40,8 @@ def test_client_recv(self):
         self.assertEqual(msg.node_id, 'identity')
 
     def test_client_retry(self):
-        server = zmqrpc.Server('0.0.0.0', 8888)
+        # use non-zero port this time, just to increase code coverage
+        server = zmqrpc.Server('127.0.0.1', 8888)
         server.socket.close()
         with self.assertRaises(zmq.error.ZMQError):
             server.recv_from_client()

setup.py

@@ -50,6 +50,7 @@
         "six>=1.10.0", 
         "pyzmq>=16.0.2", 
         "geventhttpclient-wheels==1.3.1.dev2",
+        "ConfigArgParse==0.15.1",
     ],
     test_suite="locust.test",
     tests_require=['mock'],