diff --git a/locust/runners.py b/locust/runners.py index 8286e10e83..e0e6e69d8c 100644 --- a/locust/runners.py +++ b/locust/runners.py @@ -432,6 +432,8 @@ def register_message(self, msg_type: str, listener: Callable) -> None: :param msg_type: The type of the message to listen for :param listener: The function to execute when the message is received """ + if msg_type in self.custom_messages: + raise Exception(f"Tried to register listener method for {msg_type}, but it already had a listener!") self.custom_messages[msg_type] = listener diff --git a/locust/test/test_runners.py b/locust/test/test_runners.py index f5b76b19d6..faedfe0cc2 100644 --- a/locust/test/test_runners.py +++ b/locust/test/test_runners.py @@ -626,6 +626,20 @@ def on_custom_msg(msg, **kw): msg = self.mocked_log.warning[0] self.assertIn("Unknown message type received", msg) + def test_duplicate_message_handler_registration(self): + class MyUser(User): + @task + def my_task(self): + pass + + def on_custom_msg(msg, **kw): + pass + + environment = Environment(user_classes=[MyUser]) + runner = LocalRunner(environment) + runner.register_message("test_custom_msg", on_custom_msg) + self.assertRaises(Exception, runner.register_message, "test_custom_msg", on_custom_msg) + def test_swarm_endpoint_is_non_blocking(self): class TestUser1(User): @task