From 0a0e5767976b75314443e5f82062bd31feef86d5 Mon Sep 17 00:00:00 2001
From: Andrew Baldwin <abald.cg.ide@aldogroup.com>
Date: Tue, 5 Sep 2023 19:43:57 +0200
Subject: [PATCH] Add multiple host tests for test_stats

---
 locust/test/test_stats.py | 98 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 98 insertions(+)

diff --git a/locust/test/test_stats.py b/locust/test/test_stats.py
index bdc808c4a1..1c9823249b 100644
--- a/locust/test/test_stats.py
+++ b/locust/test/test_stats.py
@@ -283,6 +283,21 @@ def test_error_grouping(self):
         self.stats.log_error("GET", "/some-path", "", Exception("Third exception!"))
         self.assertEqual(3, len(self.stats.errors))
 
+    def test_error_grouping_with_multiple_hosts(self):
+        # reset stats
+        self.stats = RequestStats()
+
+        self.stats.log_error("GET", "/some-path", "http://127.0.0.1:1", Exception("Exception!"))
+        self.stats.log_error("GET", "/some-path", "http://127.0.0.1:2", Exception("Exception!"))
+        self.assertEqual(2, len(self.stats.errors))
+        self.assertEqual(1, list(self.stats.errors.values())[0].occurrences)
+        self.assertEqual(1, list(self.stats.errors.values())[1].occurrences)
+
+        self.stats.log_error("GET", "/some-path", "http://127.0.0.1:1", Exception("Another exception!"))
+        self.stats.log_error("GET", "/some-path", "http://127.0.0.1:1", Exception("Another exception!"))
+        self.stats.log_error("GET", "/some-path", "http://127.0.0.1:2", Exception("Third exception!"))
+        self.assertEqual(4, len(self.stats.errors))
+
     def test_error_grouping_errors_with_memory_addresses(self):
         # reset stats
         self.stats = RequestStats()
@@ -373,6 +388,72 @@ def test_print_error_report(self):
         self.assertEqual(info[2], info[-2])
 
 
+class TestStatsPrintingMultipleHosts(LocustTestCase):
+    def setUp(self):
+        super().setUp()
+
+        self.stats = RequestStats()
+        for i in range(10):
+            for method, name, freq in [
+                (
+                    "GET",
+                    "test_entry",
+                    5,
+                ),
+                (
+                    "DELETE",
+                    "test_error",
+                    3,
+                ),
+            ]:
+                self.stats.log_request(method, name, f"http://127.0.0.1/{i}/{name}", i, 2000 + i)
+                if i % freq == 0:
+                    self.stats.log_error(
+                        method, name, f"http://127.0.0.1/{(i)}/{name}", RuntimeError(f"{method} error")
+                    )
+
+    def test_print_percentile_stats_multiple_hosts(self):
+        locust.stats.print_percentile_stats(self.stats)
+        info = self.mocked_log.info
+        self.assertEqual(26, len(info))
+        self.assertEqual("Response time percentiles (approximated)", info[0])
+        # check that headline contains same number of column as the value rows
+        headlines = info[1].replace("# reqs", "#reqs").split()
+        self.assertEqual(len(headlines), len(info[3].split()))
+        self.assertEqual(len(headlines) - 2, len(info[-2].split()))  # Aggregated, no "Type"
+        self.assertEqual(info[2], info[-3])  # table ascii separators
+        self.assertIn("Host", headlines)
+
+    def test_print_stats_multiple_hosts(self):
+        locust.stats.print_stats(self.stats)
+        info = self.mocked_log.info
+        self.assertEqual(25, len(info))
+
+        headlines = info[0].replace("# ", "#").split()
+        # check number of columns in separator
+        self.assertEqual(len(headlines), len(info[1].split("|")) + 2)
+        # check entry row
+        self.assertEqual(len(headlines), len(info[2].split()))
+        # check aggregated row, which is missing value in "type"-column
+        self.assertEqual(len(headlines) - 2, len(info[-2].split()))
+        # table ascii separators
+        self.assertEqual(info[1], info[-3])
+        self.assertIn("Host", headlines)
+
+    def test_print_error_report_multiple_hosts(self):
+        locust.stats.print_error_report(self.stats)
+        info = self.mocked_log.info
+        self.assertEqual(11, len(info))
+        self.assertEqual("Error report", info[0])
+
+        headlines = info[1].replace("# ", "#").split()
+        # check number of columns in headlines vs table ascii separator
+        self.assertEqual(len(headlines), len(info[2].split("|")))
+        # table ascii seprators
+        self.assertEqual(info[2], info[-2])
+        self.assertIn("Host", headlines)
+
+
 class TestCsvStats(LocustTestCase):
     STATS_BASE_NAME = "test"
     STATS_FILENAME = f"{STATS_BASE_NAME}_stats.csv"
@@ -599,6 +680,23 @@ def test_stats_history(self):
         self.assertEqual(1, len(hs1))
         self.assertEqual(0, len(hs2))
 
+    def test_csv_multiple_hosts(self):
+        class UserOne(User):
+            host = "http://127.0.0.1:1"
+
+        class UserTwo(User):
+            host = "http://127.0.0.1:2"
+
+        self.environment.user_classes = [UserOne, UserTwo]
+
+        _write_csv_files(self.environment, self.STATS_BASE_NAME)
+
+        with open(self.STATS_HISTORY_FILENAME) as f:
+            reader = csv.DictReader(f)
+            rows = [r for r in reader]
+
+        self.assertIn("Host", rows[0])
+
 
 class TestStatsEntryResponseTimesCache(unittest.TestCase):
     def setUp(self, *args, **kwargs):