From f0c6faa2dc5b6d1bc71c423dff99e8a62cf2cc02 Mon Sep 17 00:00:00 2001
From: Ankit Mehta <ankit.mehta@appian.com>
Date: Tue, 3 Dec 2019 14:22:48 -0500
Subject: [PATCH] Make csv_history file available for download on web UI

---
 locust/stats.py             | 12 ++++++++++--
 locust/templates/index.html |  2 +-
 locust/test/test_web.py     |  5 +++++
 locust/web.py               | 11 ++++++++++-
 4 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/locust/stats.py b/locust/stats.py
index 200bbc2694..248d0460b8 100644
--- a/locust/stats.py
+++ b/locust/stats.py
@@ -862,11 +862,19 @@ def stats_history_csv_header():
         '"100%"'
     )) + '\n'
 
-def stats_history_csv(stats_history_enabled=False):
+def stats_history_csv(stats_history_enabled=False, csv_for_web_ui=False):
     """Returns the Aggregated stats entry every interval"""
     from . import runners
 
-    rows = []
+    # csv_for_web_ui boolean returns the header along with the stats history row so that
+    # it can be returned as a csv for download on the web ui. Otherwise when run with
+    # the '--no-web' option we write the header first and then append the file with stats
+    # entries every interval.
+    if csv_for_web_ui:
+        rows = [stats_history_csv_header()]
+    else:
+        rows = []
+
     timestamp = int(time.time())
     stats_entries_per_iteration = []
 
diff --git a/locust/templates/index.html b/locust/templates/index.html
index a19bfa68ff..a394221890 100644
--- a/locust/templates/index.html
+++ b/locust/templates/index.html
@@ -157,7 +157,7 @@ <h2>Change the locust count</h2>
                 <div style="display:none;">
                     <div style="margin-top:20px;">
                         <a href="./stats/requests/csv">Download request statistics CSV</a><br>
-                        <a href="./stats/distribution/csv">Download response time distribution CSV</a><br>
+                        <a href="./stats/stats_history/csv">Download response time stats history CSV</a><br>
                         <a href="./stats/failures/csv">Download failures CSV</a><br>
                         <a href="./exceptions/csv">Download exceptions CSV</a>
                     </div>
diff --git a/locust/test/test_web.py b/locust/test/test_web.py
index 0037bff3b7..1fcbe23cc6 100644
--- a/locust/test/test_web.py
+++ b/locust/test/test_web.py
@@ -96,6 +96,11 @@ def test_request_stats_csv(self):
         response = requests.get("http://127.0.0.1:%i/stats/requests/csv" % self.web_port)
         self.assertEqual(200, response.status_code)
 
+    def test_request_stats_history_csv(self):
+        stats.global_stats.log_request("GET", "/test2", 120, 5612)
+        response = requests.get("http://127.0.0.1:%i/stats/stats_history/csv" % self.web_port)
+        self.assertEqual(200, response.status_code)
+
     def test_failure_stats_csv(self):
         stats.global_stats.log_error("GET", "/", Exception("Error1337"))
         response = requests.get("http://127.0.0.1:%i/stats/failures/csv" % self.web_port)
diff --git a/locust/web.py b/locust/web.py
index ec7034b5ce..b90c8b6365 100644
--- a/locust/web.py
+++ b/locust/web.py
@@ -24,7 +24,7 @@
 
 from . import runners
 from .runners import MasterLocustRunner
-from .stats import failures_csv, median_from_dict, requests_csv, sort_stats
+from .stats import failures_csv, median_from_dict, requests_csv, sort_stats, stats_history_csv
 from .util.cache import memoize
 from .util.rounding import proper_round
 
@@ -100,6 +100,15 @@ def request_stats_csv():
     response.headers["Content-disposition"] = disposition
     return response
 
+@app.route("/stats/stats_history/csv")
+def stats_history_stats_csv():
+    response = make_response(stats_history_csv(False, True))
+    file_name = "stats_history_{0}.csv".format(time())
+    disposition = "attachment;filename={0}".format(file_name)
+    response.headers["Content-type"] = "text/csv"
+    response.headers["Content-disposition"] = disposition
+    return response
+
 @app.route("/stats/failures/csv")
 def failures_stats_csv():
     response = make_response(failures_csv())