Skip to content

Commit

Permalink
KeyboardInterrupt with load shapes (#2375)
Browse files Browse the repository at this point in the history
Catch keyboard interrupt exception (CTRL+C), log it, and end test run. Fixes #1910 

Also adds integration test for this
  • Loading branch information
JevonCowell authored Jul 31, 2023
1 parent 64987e0 commit 8e8a341
Show file tree
Hide file tree
Showing 2 changed files with 52 additions and 3 deletions.
10 changes: 7 additions & 3 deletions locust/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
45 changes: 45 additions & 0 deletions locust/test/test_main.py
Original file line number Diff line number Diff line change
Expand Up @@ -1185,6 +1185,51 @@ 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):
print("running my_task()")
"""
)
) as mocked:
proc = subprocess.Popen(
[
"locust",
"-f",
mocked,
"--headless",
],
stdout=PIPE,
stderr=PIPE,
text=True,
)
gevent.sleep(1.9)
proc.send_signal(signal.SIGINT)
stdout, stderr = proc.communicate()
print(stderr, stdout)
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
self.assertRegex(stderr, r".*Aggregated[\S\s]*Shutting down[\S\s]*Aggregated.*")


class DistributedIntegrationTests(ProcessIntegrationTest):
def test_expect_workers(self):
Expand Down

0 comments on commit 8e8a341

Please sign in to comment.