From 0df7920bce921bcb7bdb9a642e0b43eac0ae49a4 Mon Sep 17 00:00:00 2001
From: jcowell <jcowell@atlassian.com>
Date: Sun, 30 Jul 2023 19:16:47 -0400
Subject: [PATCH 1/4] ISSUE-1910: catch keyboard interrupt exception, log it,
 and end test.

---
 locust/main.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/locust/main.py b/locust/main.py
index df6caa3b99..cab7824e8b 100644
--- a/locust/main.py
+++ b/locust/main.py
@@ -378,9 +378,13 @@ def start_automatic_run():
                 if options.run_time:
                     sys.stderr.write("It makes no sense to combine --run-time and LoadShapes. Bailing out.\n")
                     sys.exit(1)
-                environment.runner.start_shape()
-                environment.runner.shape_greenlet.join()
-                stop_and_optionally_quit()
+                try:
+                    environment.runner.start_shape()
+                    environment.runner.shape_greenlet.join()
+                except KeyboardInterrupt:
+                    logging.info("Exiting due to CTRL+C interruption")
+                finally:
+                    stop_and_optionally_quit()
             else:
                 headless_master_greenlet = gevent.spawn(runner.start, options.num_users, options.spawn_rate)
                 headless_master_greenlet.link_exception(greenlet_exception_handler)

From 669f32e5ff79f6888a6ae42ed542faf7fa0ef7ea Mon Sep 17 00:00:00 2001
From: jcowell <jcowell@atlassian.com>
Date: Sun, 30 Jul 2023 19:17:15 -0400
Subject: [PATCH 2/4] ISSUE-1910: add integration test for keyboard interrupt

---
 locust/test/test_main.py | 42 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 42 insertions(+)

diff --git a/locust/test/test_main.py b/locust/test/test_main.py
index cf76cd9342..d543a42dd5 100644
--- a/locust/test/test_main.py
+++ b/locust/test/test_main.py
@@ -1185,6 +1185,48 @@ def task1(self):
             self.assertIn("No tasks defined on MyUser", stderr)
             self.assertEqual(1, proc.returncode)
 
+    def test_graceful_exit_when_keyboard_interrupt(self):
+        with temporary_file(
+            content=textwrap.dedent(
+                """
+                from locust import User, events,task, between, LoadTestShape
+                @events.test_stop.add_listener
+                def on_test_stop(environment, **kwargs) -> None:
+                    print("Test Stopped")
+
+                class LoadTestShape(LoadTestShape):
+                    def tick(self):
+                        run_time = self.get_run_time()
+                        if run_time < 2:
+                            return (10, 1)
+
+                        return None
+
+                class TestUser(User):
+                    @task
+                    def my_task(self):
+                        raise KeyboardInterrupt
+            """
+            )
+        ) as mocked:
+            proc = subprocess.Popen(
+                [
+                    "locust",
+                    "-f",
+                    mocked,
+                    "--headless",
+                ],
+                stdout=PIPE,
+                stderr=PIPE,
+                text=True,
+            )
+            gevent.sleep(1.9)
+            stdout, stderr = proc.communicate()
+            print(stderr, stdout)
+            self.assertIn("Starting Locust", stderr)
+            self.assertIn("Exiting due to CTRL+C interruption", stderr)
+            self.assertIn("Test Stopped", stdout)
+
 
 class DistributedIntegrationTests(ProcessIntegrationTest):
     def test_expect_workers(self):

From 025297058813c49226cb2c08b3fa2a604a68e761 Mon Sep 17 00:00:00 2001
From: jcowell <jcowell@atlassian.com>
Date: Sun, 30 Jul 2023 19:23:07 -0400
Subject: [PATCH 3/4] ISSUE-1910: confirm that test ends with report

---
 locust/test/test_main.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/locust/test/test_main.py b/locust/test/test_main.py
index d543a42dd5..85ec8c6707 100644
--- a/locust/test/test_main.py
+++ b/locust/test/test_main.py
@@ -1226,6 +1226,8 @@ def my_task(self):
             self.assertIn("Starting Locust", stderr)
             self.assertIn("Exiting due to CTRL+C interruption", stderr)
             self.assertIn("Test Stopped", stdout)
+            # ensure stats printer printed at least one report before shutting down and that there was a final report printed as well
+            self.assertRegex(stderr, r".*Aggregated[\S\s]*Shutting down[\S\s]*Aggregated.*")
 
 
 class DistributedIntegrationTests(ProcessIntegrationTest):

From f4b0af7a8ce0a87687ff69fa17c2eb95f35aac91 Mon Sep 17 00:00:00 2001
From: jcowell <jcowell@atlassian.com>
Date: Mon, 31 Jul 2023 10:31:59 -0400
Subject: [PATCH 4/4] ISSUE-1910: send SIGINT instead to process

---
 locust/test/test_main.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/locust/test/test_main.py b/locust/test/test_main.py
index 85ec8c6707..6048f3d15a 100644
--- a/locust/test/test_main.py
+++ b/locust/test/test_main.py
@@ -1205,7 +1205,7 @@ def tick(self):
                 class TestUser(User):
                     @task
                     def my_task(self):
-                        raise KeyboardInterrupt
+                        print("running my_task()")
             """
             )
         ) as mocked:
@@ -1221,9 +1221,10 @@ def my_task(self):
                 text=True,
             )
             gevent.sleep(1.9)
+            proc.send_signal(signal.SIGINT)
             stdout, stderr = proc.communicate()
             print(stderr, stdout)
-            self.assertIn("Starting Locust", stderr)
+            self.assertIn("Shape test starting", stderr)
             self.assertIn("Exiting due to CTRL+C interruption", stderr)
             self.assertIn("Test Stopped", stdout)
             # ensure stats printer printed at least one report before shutting down and that there was a final report printed as well