Merge pull request #1428 from ajt89/csv-output-fixes

Use csv module to generate csv data

locust/stats.py

@@ -1,3 +1,4 @@
+import csv
 import hashlib
 import time
 from collections import namedtuple, OrderedDict
@@ -767,59 +768,56 @@ def stats_writer(environment, base_filepath, full_history=False):
 def write_csv_files(environment, base_filepath, full_history=False):
     """Writes the requests, distribution, and failures csvs."""
     with open(base_filepath + '_stats.csv', 'w') as f:
-        f.write(requests_csv(environment.stats))
+        csv_writer = csv.writer(f)
+        requests_csv(environment.stats, csv_writer)
 
     with open(base_filepath + '_stats_history.csv', 'a') as f:
         f.write(stats_history_csv(environment, full_history) + "\n")
 
     with open(base_filepath + '_failures.csv', 'w') as f:
-        f.write(failures_csv(environment.stats))
+        csv_writer = csv.writer(f)
+        failures_csv(environment.stats, csv_writer)
 
 
 def sort_stats(stats):
     return [stats[key] for key in sorted(stats.keys())]
 
 
-def requests_csv(stats):
-    from . import runners
-
+def requests_csv(stats, csv_writer):
     """Returns the contents of the 'requests' & 'distribution' tab as CSV."""
-    rows = [
-        ",".join([
-            '"Type"',
-            '"Name"',
-            '"Request Count"',
-            '"Failure Count"',
-            '"Median Response Time"',
-            '"Average Response Time"',
-            '"Min Response Time"',
-            '"Max Response Time"',
-            '"Average Content Size"',
-            '"Requests/s"',
-            '"Failures/s"',
-            '"50%"',
-            '"66%"',
-            '"75%"',
-            '"80%"',
-            '"90%"',
-            '"95%"',
-            '"98%"',
-            '"99%"',
-            '"99.9%"',
-            '"99.99%"',
-            '"99.999%"',
-            '"100%"'
-        ])
-    ]
+    csv_writer.writerow([
+        "Type",
+        "Name",
+        "Request Count",
+        "Failure Count",
+        "Median Response Time",
+        "Average Response Time",
+        "Min Response Time",
+        "Max Response Time",
+        "Average Content Size",
+        "Requests/s",
+        "Failures/s",
+        "50%",
+        "66%",
+        "75%",
+        "80%",
+        "90%",
+        "95%",
+        "98%",
+        "99%",
+        "99.9%",
+        "99.99%",
+        "99.999%",
+        "100%",
+    ])
 
     for s in chain(sort_stats(stats.entries), [stats.total]):
         if s.num_requests:
-            percentile_str = ','.join([
-                str(int(s.get_response_time_percentile(x) or 0)) for x in PERCENTILES_TO_REPORT])
+            percentile_row = [int(s.get_response_time_percentile(x) or 0) for x in PERCENTILES_TO_REPORT]
         else:
-            percentile_str = ','.join(['"N/A"'] * len(PERCENTILES_TO_REPORT))
+            percentile_row = ["N/A"] * len(PERCENTILES_TO_REPORT)
 
-        rows.append('"%s","%s",%i,%i,%i,%i,%i,%i,%i,%.2f,%.2f,%s' % (
+        stats_row = [
             s.method,
             s.name,
             s.num_requests,
@@ -831,10 +829,9 @@ def requests_csv(stats):
             s.avg_content_length,
             s.total_rps,
             s.total_fail_per_sec,
-            percentile_str
-        ))
-    return "\n".join(rows)
+        ]
 
+        csv_writer.writerow(stats_row + percentile_row)
 
 def stats_history_csv_header():
     """Headers for the stats history CSV"""
@@ -906,22 +903,19 @@ def stats_history_csv(environment, all_entries=False):
 
     return "\n".join(rows)
 
-def failures_csv(stats):
+def failures_csv(stats, csv_writer):
     """"Return the contents of the 'failures' tab as a CSV."""
-    rows = [
-        ",".join((
-            '"Method"',
-            '"Name"',
-            '"Error"',
-            '"Occurrences"',
-        ))
-    ]
+    csv_writer.writerow([
+        "Method",
+        "Name",
+        "Error",
+        "Occurrences",
+    ])
 
     for s in sort_stats(stats.errors):
-        rows.append('"%s","%s","%s",%i' % (
+        csv_writer.writerow([
             s.method,
             s.name,
             s.error,
             s.occurrences,
-        ))
-    return "\n".join(rows)
+        ])

locust/test/test_stats.py

@@ -3,6 +3,7 @@
 import unittest
 import re
 import os
+import json
 
 import gevent
 import mock
@@ -415,6 +416,32 @@ def t(self):
             self.assertEqual("%i" % (i + 1), row["Total Request Count"])
             self.assertGreaterEqual(int(row["Timestamp"]), start_time)
 
+    def test_requests_csv_quote_escaping(self):
+        with mock.patch("locust.rpc.rpc.Server", mocked_rpc()) as server:
+            environment = Environment()
+            master = environment.create_master_runner(master_bind_host="*", master_bind_port=0)
+            server.mocked_send(Message("client_ready", None, "fake_client"))
+
+            request_name_dict = {
+                "scenario": "get cashes",
+                "path": "/cash/[amount]",
+                "arguments": [{"size": 1}],
+            }
+            request_name_str = json.dumps(request_name_dict)
+
+            master.stats.get(request_name_str, "GET").log(100, 23455)
+            data = {"user_count": 1}
+            environment.events.report_to_master.fire(client_id="fake_client", data=data)
+            master.stats.clear_all()
+            server.mocked_send(Message("stats", data, "fake_client"))
+
+            locust.stats.write_csv_files(environment, self.STATS_BASE_NAME, full_history=True)
+            with open(self.STATS_FILENAME) as f:
+                reader = csv.DictReader(f)
+                rows = [r for r in reader]
+                csv_request_name = rows[0].get("Name")
+                self.assertEqual(request_name_str, csv_request_name)
+
 
 class TestStatsEntryResponseTimesCache(unittest.TestCase):
     def setUp(self, *args, **kwargs):

locust/web.py

@@ -172,7 +172,10 @@ def reset_stats():
         @app.route("/stats/requests/csv")
         @self.auth_required_if_enabled
         def request_stats_csv():
-            response = make_response(requests_csv(self.environment.runner.stats))
+            data = StringIO()
+            writer = csv.writer(data)
+            requests_csv(self.environment.runner.stats, writer)
+            response = make_response(data.getvalue())
             file_name = "requests_{0}.csv".format(time())
             disposition = "attachment;filename={0}".format(file_name)
             response.headers["Content-type"] = "text/csv"
@@ -182,7 +185,10 @@ def request_stats_csv():
         @app.route("/stats/failures/csv")
         @self.auth_required_if_enabled
         def failures_stats_csv():
-            response = make_response(failures_csv(self.environment.runner.stats))
+            data = StringIO()
+            writer = csv.writer(data)
+            failures_csv(self.environment.runner.stats, writer)
+            response = make_response(data.getvalue())
             file_name = "failures_{0}.csv".format(time())
             disposition = "attachment;filename={0}".format(file_name)
             response.headers["Content-type"] = "text/csv"
@@ -267,9 +273,8 @@ def exceptions_csv():
             for exc in environment.runner.exceptions.values():
                 nodes = ", ".join(exc["nodes"])
                 writer.writerow([exc["count"], exc["msg"], exc["traceback"], nodes])
-
-            data.seek(0)
-            response = make_response(data.read())
+
+            response = make_response(data.getvalue())
             file_name = "exceptions_{0}.csv".format(time())
             disposition = "attachment;filename={0}".format(file_name)
             response.headers["Content-type"] = "text/csv"