Merge pull request #2481 from sillydan1/feature/create-csv-prefix-dirs

Create any directories as part of the CSV Prefix
locustio · Nov 23, 2023 · 04f9a66 · 04f9a66
2 parents 8472ed8 + edc3424
commit 04f9a66
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 2 deletions.
diff --git a/locust/argument_parser.py b/locust/argument_parser.py
@@ -500,7 +500,7 @@ def setup_parser_arguments(parser):
         "--csv",  # Name repeated in 'parse_options'
         dest="csv_prefix",
         metavar="<filename>",
-        help="Store request stats to files in CSV format. Setting this option will generate three files: <filename>_stats.csv, <filename>_stats_history.csv and <filename>_failures.csv",
+        help="Store request stats to files in CSV format. Setting this option will generate three files: <filename>_stats.csv, <filename>_stats_history.csv and <filename>_failures.csv. Any folders part of the prefix will be automatically created",
         env_var="LOCUST_CSV",
     )
     stats_group.add_argument(

diff --git a/locust/main.py b/locust/main.py
@@ -398,6 +398,10 @@ def kill_workers(children):
             sys.exit(1)
 
     if options.csv_prefix:
+        base_csv_file = os.path.basename(options.csv_prefix)
+        base_csv_dir = options.csv_prefix[: -len(base_csv_file)]
+        if not os.path.exists(base_csv_dir):
+            os.makedirs(base_csv_dir)
         stats_csv_writer = StatsCSVFileWriter(
             environment, stats.PERCENTILES_TO_REPORT, options.csv_prefix, options.stats_history_enabled
         )

diff --git a/locust/test/test_web.py b/locust/test/test_web.py
@@ -1102,6 +1102,7 @@ def test_index_with_https(self):
 
 
 class TestWebUIFullHistory(LocustTestCase, _HeaderCheckMixin):
+    STATS_BASE_DIR = "csv_output"
     STATS_BASE_NAME = "web_test"
     STATS_FILENAME = f"{STATS_BASE_NAME}_stats.csv"
     STATS_HISTORY_FILENAME = f"{STATS_BASE_NAME}_stats_history.csv"
@@ -1112,7 +1113,9 @@ def setUp(self):
         self.remove_files_if_exists()
 
         parser = get_parser(default_config_files=[])
-        self.environment.parsed_options = parser.parse_args(["--csv", self.STATS_BASE_NAME, "--csv-full-history"])
+        self.environment.parsed_options = parser.parse_args(
+            ["--csv", os.path.join(self.STATS_BASE_DIR, self.STATS_BASE_NAME), "--csv-full-history"]
+        )
         self.stats = self.environment.stats
         self.stats.CSV_STATS_INTERVAL_SEC = 0.02
 
@@ -1139,6 +1142,7 @@ def remove_files_if_exists(self):
         self.remove_file_if_exists(self.STATS_FILENAME)
         self.remove_file_if_exists(self.STATS_HISTORY_FILENAME)
         self.remove_file_if_exists(self.STATS_FAILURES_FILENAME)
+        self.remove_file_if_exists(self.STATS_BASE_DIR)
 
     def test_request_stats_full_history_csv(self):
         self.stats.log_request("GET", "/test", 1.39764125, 2)