Skip to content

Commit

Permalink
Throw an exception if someone tries to register a custom message hand…
Browse files Browse the repository at this point in the history
…ler twice (it is kind of undefined what should happen in that case, so we just dont!)
  • Loading branch information
cyberw committed Dec 12, 2023
1 parent 247fee9 commit 3bb7826
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 0 deletions.
2 changes: 2 additions & 0 deletions locust/runners.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down
14 changes: 14 additions & 0 deletions locust/test/test_runners.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 3bb7826

Please sign in to comment.