locust 2.23.2.dev26__py3-none-any.whl → 2.23.2.dev39__py3-none-any.whl

locust/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '2.23.2.dev26'
-__version_tuple__ = version_tuple = (2, 23, 2, 'dev26')
+__version__ = version = '2.23.2.dev39'
+__version_tuple__ = version_tuple = (2, 23, 2, 'dev39')

locust/argument_parser.py

@@ -12,10 +12,16 @@ import socket
 import sys
 import tempfile
 import textwrap
+from collections import OrderedDict
 from typing import Any, NamedTuple
 from urllib.parse import urlparse
 from uuid import uuid4
 
+if sys.version_info >= (3, 11):
+    import tomllib
+else:
+    import tomli as tomllib
+
 import configargparse
 import gevent
 import requests
@@ -23,7 +29,7 @@ import requests
 version = locust.__version__
 
 
-DEFAULT_CONFIG_FILES = ["~/.locust.conf", "locust.conf"]
+DEFAULT_CONFIG_FILES = ("~/.locust.conf", "locust.conf", "pyproject.toml")
 
 
 class LocustArgumentParser(configargparse.ArgumentParser):
@@ -63,6 +69,31 @@ class LocustArgumentParser(configargparse.ArgumentParser):
         }
 
 
+class LocustTomlConfigParser(configargparse.TomlConfigParser):
+    def parse(self, stream):
+        try:
+            config = tomllib.loads(stream.read())
+        except Exception as e:
+            raise configargparse.ConfigFileParserException(f"Couldn't parse TOML file: {e}")
+
+        # convert to dict and filter based on section names
+        result = OrderedDict()
+
+        for section in self.sections:
+            data = configargparse.get_toml_section(config, section)
+            if data:
+                for key, value in data.items():
+                    if isinstance(value, list):
+                        result[key] = value
+                    elif value is None:
+                        pass
+                    else:
+                        result[key] = str(value)
+                break
+
+        return result
+
+
 def _is_package(path):
     """
     Is the given path a Python package?
@@ -186,6 +217,12 @@ def download_locustfile_from_url(url: str) -> str:
 def get_empty_argument_parser(add_help=True, default_config_files=DEFAULT_CONFIG_FILES) -> LocustArgumentParser:
     parser = LocustArgumentParser(
         default_config_files=default_config_files,
+        config_file_parser_class=configargparse.CompositeConfigParser(
+            [
+                LocustTomlConfigParser(["tool.locust"]),
+                configargparse.DefaultConfigFileParser,
+            ]
+        ),
         add_env_var_help=False,
         add_config_file_help=False,
         add_help=add_help,

locust/test/test_parser.py

@@ -38,17 +38,50 @@ class TestParser(unittest.TestCase):
         opts = self.parser.parse_args(args)
         self.assertEqual(opts.skip_log_setup, True)
 
-    def test_parameter_parsing(self):
-        with NamedTemporaryFile(mode="w") as file:
-            os.environ["LOCUST_LOCUSTFILE"] = "locustfile_from_env"
-            file.write("host host_from_config\nweb-host webhost_from_config")
+    def test_parse_options_from_conf_file(self):
+        with NamedTemporaryFile(mode="w", suffix=".conf") as file:
+            config_data = """\
+            locustfile = ./test_locustfile.py
+            web-host = 127.0.0.1
+            web-port = 45787
+            headless
+            tags = [Critical, Normal]
+            """
+
+            file.write(config_data)
             file.flush()
             parser = get_parser(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
+            options = parser.parse_args(["-H", "https://example.com"])
+
+        self.assertEqual("./test_locustfile.py", options.locustfile)
+        self.assertEqual("127.0.0.1", options.web_host)
+        self.assertEqual(45787, options.web_port)
+        self.assertTrue(options.headless)
+        self.assertEqual(["Critical", "Normal"], options.tags)
+        self.assertEqual("https://example.com", options.host)
+
+    def test_parse_options_from_toml_file(self):
+        with NamedTemporaryFile(mode="w", suffix=".toml") as file:
+            config_data = """\
+            [tool.locust]
+            locustfile = "./test_locustfile.py"
+            web-host = "127.0.0.1"
+            web-port = 45787
+            headless = true
+            tags = ["Critical", "Normal"]
+            """
+
+            file.write(config_data)
+            file.flush()
+            parser = get_parser(default_config_files=[file.name])
+            options = parser.parse_args(["-H", "https://example.com"])
+
+        self.assertEqual("./test_locustfile.py", options.locustfile)
+        self.assertEqual("127.0.0.1", options.web_host)
+        self.assertEqual(45787, options.web_port)
+        self.assertTrue(options.headless)
+        self.assertEqual(["Critical", "Normal"], options.tags)
+        self.assertEqual("https://example.com", options.host)
 
 
 class TestArgumentParser(LocustTestCase):