Merge pull request #1316 from locustio/csv-stats-improvements

Add User count to CSV history stats

docs/changelog.rst

@@ -27,6 +27,8 @@ Breaking changes
       def on_test_start(**kw):
           print("test is stopping")
 * ``TaskSequence`` and ``@seq_task`` has been replaced with :ref:`SequentialTaskSet <sequential-taskset>`.
+* A ``User count`` column has been added to the history stats CSV file, and the column order has been changed 
+  (``Timestamp`` is now the first column).
 
 
 0.14.0

locust/main.py

@@ -260,7 +260,7 @@ def timelimit_stop():
         stats_printer_greenlet = gevent.spawn(stats_printer(runner.stats))
 
     if options.csvfilebase:
-        gevent.spawn(stats_writer, runner.stats, options.csvfilebase, full_history=options.stats_history_enabled)
+        gevent.spawn(stats_writer, environment, options.csvfilebase, full_history=options.stats_history_enabled)
 
 
     def shutdown(code=0):
@@ -278,7 +278,7 @@ def shutdown(code=0):
         print_stats(runner.stats, current=False)
         print_percentile_stats(runner.stats)
         if options.csvfilebase:
-            write_csv_files(runner.stats, options.csvfilebase, options.stats_history_enabled)
+            write_csv_files(environment, options.csvfilebase, full_history=options.stats_history_enabled)
         print_error_report(runner.stats)
         sys.exit(code)
 

locust/stats.py

@@ -751,25 +751,25 @@ def stats_printer_func():
             gevent.sleep(CONSOLE_STATS_INTERVAL_SEC)
     return stats_printer_func
 
-def stats_writer(stats, base_filepath, full_history=False):
+def stats_writer(environment, base_filepath, full_history=False):
     """Writes the csv files for the locust run."""
     with open(base_filepath + '_stats_history.csv', 'w') as f:
         f.write(stats_history_csv_header())
     while True:
-        write_csv_files(stats, base_filepath, full_history)
+        write_csv_files(environment, base_filepath, full_history)
         gevent.sleep(CSV_STATS_INTERVAL_SEC)
 
 
-def write_csv_files(stats, 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(stats))
+        f.write(requests_csv(environment.runner.stats))
 
     with open(base_filepath + '_stats_history.csv', 'a') as f:
-        f.write(stats_history_csv(stats, full_history) + "\n")
+        f.write(stats_history_csv(environment, full_history) + "\n")
 
     with open(base_filepath + '_failures.csv', 'w') as f:
-        f.write(failures_csv(stats))
+        f.write(failures_csv(environment.runner.stats))
 
 
 def sort_stats(stats):
@@ -831,13 +831,15 @@ def requests_csv(stats):
         ))
     return "\n".join(rows)
 
+
 def stats_history_csv_header():
     """Headers for the stats history CSV"""
 
     return ','.join((
+        '"Timestamp"',
+        '"User count"',
         '"Type"',
         '"Name"',
-        '"Timestamp"',
         '"# requests"',
         '"# failures"',
         '"Requests/s"',
@@ -861,34 +863,31 @@ def stats_history_csv_header():
         '"100%"'
     )) + '\n'
 
-def stats_history_csv(stats, stats_history_enabled=False, csv_for_web_ui=False):
-    """Returns the Aggregated stats entry every interval"""
-    # 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 '--headless' 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 = []
-
+def stats_history_csv(environment, all_entries=False):
+    """
+    Return a string of CSV rows with the *current* stats. By default only includes the 
+    Aggregated stats entry, but if all_entries is set to True, a row for each entry will 
+    will be included.
+    """
+    stats = environment.runner.stats
     timestamp = int(time.time())
-    stats_entries_per_iteration = []
-
-    if stats_history_enabled:
-        stats_entries_per_iteration = sort_stats(stats.entries)
-
-    for s in chain(stats_entries_per_iteration, [stats.total]):
+    stats_entries = []
+    if all_entries:
+        stats_entries = sort_stats(stats.entries)
+
+    rows = []
+    for s in chain(stats_entries, [stats.total]):
         if s.num_requests:
             percentile_str = ','.join([
                 str(int(s.get_current_response_time_percentile(x) or 0)) for x in PERCENTILES_TO_REPORT])
         else:
             percentile_str = ','.join(['"N/A"'] * len(PERCENTILES_TO_REPORT))
 
-        rows.append('"%s","%s","%s",%i,%i,%.2f,%.2f,%i,%i,%i,%.2f,%.2f,%s' % (
+        rows.append('"%i","%i","%s","%s",%i,%i,%.2f,%.2f,%i,%i,%i,%.2f,%.2f,%s' % (
+            timestamp,
+            environment.runner.user_count,
             s.method,
             s.name,
-            timestamp,
             s.num_requests,
             s.num_failures,
             s.current_rps,

locust/templates/index.html

@@ -169,7 +169,6 @@ <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/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>

locust/test/test_stats.py

@@ -10,10 +10,11 @@
 from locust.core import HttpLocust, TaskSet, task, Locust
 from locust.env import Environment
 from locust.inspectlocust import get_task_ratio_dict
-from locust.runners import MasterLocustRunner
+from locust.runners import LocalLocustRunner, MasterLocustRunner
 from locust.rpc.protocol import Message
 from locust.stats import CachedResponseTimes, RequestStats, StatsEntry, diff_response_time_dicts, stats_writer
 from locust.test.testcases import LocustTestCase
+from locust.wait_time import constant
 
 from .testcases import WebserverTestCase
 from .test_runners import mocked_options, mocked_rpc
@@ -286,31 +287,19 @@ def test_print_percentile_stats(self):
         self.assertEqual(len(headlines), len(info[5].split()))
 
 
-class TestWriteStatCSVs(LocustTestCase):
+class TestCsvStats(LocustTestCase):
     STATS_BASE_NAME = "test"
     STATS_FILENAME = "{}_stats.csv".format(STATS_BASE_NAME)
     STATS_HISTORY_FILENAME = "{}_stats_history.csv".format(STATS_BASE_NAME)
     STATS_FAILURES_FILENAME = "{}_failures.csv".format(STATS_BASE_NAME)
 
     def setUp(self):
         super().setUp()
-        class User(Locust):
-            setup_run_count = 0
-            task_run_count = 0
-            locust_error_count = 0
-            wait_time = locust.wait_time.constant(1)
-            @task
-            def my_task(self):
-                User.task_run_count += 1
-
-        self.environment = Environment(options=mocked_options())
-        locust.runners.locust_runner = locust.runners.LocalLocustRunner(self.environment, [User])
         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)
 
     def tearDown(self):
-        locust.runners.locust_runner.quit()
         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)
@@ -320,20 +309,20 @@ def remove_file_if_exists(self, filename):
             os.remove(filename)
 
     def test_write_csv_files(self):
-        locust.stats.write_csv_files(self.runner.stats, self.STATS_BASE_NAME)
+        locust.stats.write_csv_files(self.environment, self.STATS_BASE_NAME)
         self.assertTrue(os.path.exists(self.STATS_FILENAME))
         self.assertTrue(os.path.exists(self.STATS_HISTORY_FILENAME))
         self.assertTrue(os.path.exists(self.STATS_FAILURES_FILENAME))
 
     def test_write_csv_files_full_history(self):
-        locust.stats.write_csv_files(self.runner.stats, self.STATS_BASE_NAME, full_history=True)
+        locust.stats.write_csv_files(self.environment, self.STATS_BASE_NAME, full_history=True)
         self.assertTrue(os.path.exists(self.STATS_FILENAME))
         self.assertTrue(os.path.exists(self.STATS_HISTORY_FILENAME))
         self.assertTrue(os.path.exists(self.STATS_FAILURES_FILENAME))
 
     @mock.patch("locust.stats.CSV_STATS_INTERVAL_SEC", new=0.2)
     def test_csv_stats_writer(self):
-        greenlet = gevent.spawn(stats_writer, self.runner.stats, self.STATS_BASE_NAME)
+        greenlet = gevent.spawn(stats_writer, self.environment, self.STATS_BASE_NAME)
         gevent.sleep(0.21)
         gevent.kill(greenlet)
         self.assertTrue(os.path.exists(self.STATS_FILENAME))
@@ -351,7 +340,7 @@ def test_csv_stats_writer(self):
     @mock.patch("locust.stats.CSV_STATS_INTERVAL_SEC", new=0.2)
     def test_csv_stats_writer_full_history(self):
         self.runner.stats.log_request("GET", "/", 10, content_length=666)
-        greenlet = gevent.spawn(stats_writer, self.runner.stats, self.STATS_BASE_NAME, full_history=True)
+        greenlet = gevent.spawn(stats_writer, self.environment, self.STATS_BASE_NAME, full_history=True)
         gevent.sleep(0.21)
         gevent.kill(greenlet)
         self.assertTrue(os.path.exists(self.STATS_FILENAME))
@@ -386,10 +375,45 @@ def test_csv_stats_on_master_from_aggregated_stats(self):
             s = master.stats.get("/", "GET")
             self.assertEqual(700, s.median_response_time)
 
-            locust.stats.write_csv_files(master.stats, self.STATS_BASE_NAME, full_history=True)
+            locust.stats.write_csv_files(self.environment, self.STATS_BASE_NAME, full_history=True)
             self.assertTrue(os.path.exists(self.STATS_FILENAME))
             self.assertTrue(os.path.exists(self.STATS_HISTORY_FILENAME))
             self.assertTrue(os.path.exists(self.STATS_FAILURES_FILENAME))
+
+    @mock.patch("locust.stats.CSV_STATS_INTERVAL_SEC", new=0.2)
+    def test_user_count_in_csv_history_stats(self):
+        start_time = int(time.time())
+        class TestUser(Locust):
+            wait_time = constant(10)
+            @task
+            def t(self):
+                self.environment.runner.stats.log_request("GET", "/", 10, 10)
+        runner = LocalLocustRunner(self.environment, [TestUser])
+        runner.start(3, 5) # spawn a user every 0.2 second
+        gevent.sleep(0.1)
+
+        greenlet = gevent.spawn(stats_writer, self.environment, self.STATS_BASE_NAME, full_history=True)
+        gevent.sleep(0.6)
+        gevent.kill(greenlet)
+
+        runner.stop()
+
+        with open(self.STATS_HISTORY_FILENAME) as f:
+            reader = csv.DictReader(f)
+            rows = [r for r in reader]
+
+        self.assertEqual(6, len(rows))
+        for i in range(3):
+            row = rows.pop(0)
+            self.assertEqual("%i" % (i + 1), row["User count"])
+            self.assertEqual("/", row["Name"])
+            self.assertEqual("%i" % (i + 1), row["# requests"])
+            self.assertGreaterEqual(int(row["Timestamp"]), start_time)
+            row = rows.pop(0)
+            self.assertEqual("%i" % (i + 1), row["User count"])
+            self.assertEqual("Aggregated", row["Name"])
+            self.assertEqual("%i" % (i + 1), row["# requests"])
+            self.assertGreaterEqual(int(row["Timestamp"]), start_time)
 
 
 class TestStatsEntryResponseTimesCache(unittest.TestCase):

locust/test/test_web.py

@@ -108,11 +108,6 @@ 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):
-        self.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):
         self.stats.log_error("GET", "/", Exception("Error1337"))
         response = requests.get("http://127.0.0.1:%i/stats/failures/csv" % self.web_port)

locust/web.py

@@ -23,7 +23,7 @@
 
 from . import runners
 from .runners import MasterLocustRunner
-from .stats import failures_csv, median_from_dict, requests_csv, sort_stats, stats_history_csv
+from .stats import failures_csv, median_from_dict, requests_csv, sort_stats
 from .util.cache import memoize
 from .util.rounding import proper_round
 from .util.timespan import parse_timespan
@@ -120,15 +120,6 @@ 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(self.environment.runner.stats, 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(self.environment.runner.stats))