From 2187721ec63aabf91fb64b2f73a9d76729b644f8 Mon Sep 17 00:00:00 2001 From: ajt89 Date: Sun, 14 Jun 2020 23:35:22 -0700 Subject: [PATCH 1/2] Use csv module to generate csv data --- locust/stats.py | 96 ++++++++++++++++++--------------------- locust/test/test_stats.py | 27 +++++++++++ locust/web.py | 12 ++++- 3 files changed, 82 insertions(+), 53 deletions(-) diff --git a/locust/stats.py b/locust/stats.py index 9cc8febb6c..19e273498d 100644 --- a/locust/stats.py +++ b/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) + ]) diff --git a/locust/test/test_stats.py b/locust/test/test_stats.py index f5eaee59e8..75787e69eb 100644 --- a/locust/test/test_stats.py +++ b/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): diff --git a/locust/web.py b/locust/web.py index 2af8022545..7b0aec481a 100644 --- a/locust/web.py +++ b/locust/web.py @@ -172,7 +172,11 @@ 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) + data.seek(0) + response = make_response(data.read()) file_name = "requests_{0}.csv".format(time()) disposition = "attachment;filename={0}".format(file_name) response.headers["Content-type"] = "text/csv" @@ -182,7 +186,11 @@ 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) + data.seek(0) + response = make_response(data.read()) file_name = "failures_{0}.csv".format(time()) disposition = "attachment;filename={0}".format(file_name) response.headers["Content-type"] = "text/csv" From 494e4d913c9aefec66f10a6d9a9d6808bf065932 Mon Sep 17 00:00:00 2001 From: ajt89 Date: Mon, 15 Jun 2020 14:01:30 -0700 Subject: [PATCH 2/2] Use getvalue to serve csv web responses --- locust/web.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/locust/web.py b/locust/web.py index 7b0aec481a..2c1aabdf13 100644 --- a/locust/web.py +++ b/locust/web.py @@ -175,8 +175,7 @@ def request_stats_csv(): data = StringIO() writer = csv.writer(data) requests_csv(self.environment.runner.stats, writer) - data.seek(0) - response = make_response(data.read()) + 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" @@ -189,8 +188,7 @@ def failures_stats_csv(): data = StringIO() writer = csv.writer(data) failures_csv(self.environment.runner.stats, writer) - data.seek(0) - response = make_response(data.read()) + 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" @@ -275,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"