From 089d525defbfec50d34fa334684f7c9d688543c8 Mon Sep 17 00:00:00 2001
From: Nocturn9x <hackhab@gmail.com>
Date: Wed, 10 May 2023 12:13:47 +0200
Subject: [PATCH] Minor addition to chatroom test

---
 tests/chatroom_server.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tests/chatroom_server.py b/tests/chatroom_server.py
index 33be043..5ba9f77 100644
--- a/tests/chatroom_server.py
+++ b/tests/chatroom_server.py
@@ -20,6 +20,10 @@ async def message_handler(q: aiosched.Queue):
             logging.info(f"Got message {msg!r} with the following payload: {payload}")
     except Exception as e:
         logging.error(f"An exception occurred in the message handler -> {type(e).__name__}: {e}")
+    except aiosched.errors.Cancelled:
+        logging.warning(f"Cancellation detected, message handler shutting down")
+        # Propagate the cancellation
+        raise
 
 
 async def serve(bind_address: tuple):