From 6b664c26fbedfe570403d6907878150ad76a7b97 Mon Sep 17 00:00:00 2001
From: krassowski <5832902+krassowski@users.noreply.github.com>
Date: Fri, 12 Jul 2024 17:46:26 +0100
Subject: [PATCH] Add a test for not sending closing message twice

---
 .../jupyter_ai/tests/test_handlers.py         | 24 +++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/packages/jupyter-ai/jupyter_ai/tests/test_handlers.py b/packages/jupyter-ai/jupyter_ai/tests/test_handlers.py
index a38ea6a26..6fab65763 100644
--- a/packages/jupyter-ai/jupyter_ai/tests/test_handlers.py
+++ b/packages/jupyter-ai/jupyter_ai/tests/test_handlers.py
@@ -133,6 +133,30 @@ async def test_default_closes_pending_on_error(human_chat_message):
     assert isinstance(handler.messages[1], ClosePendingMessage)
 
 
+async def test_sends_closing_message_at_most_once(human_chat_message):
+    handler = TestDefaultChatHandler(
+        lm_provider=MockProvider,
+        lm_provider_params={
+            "model_id": "model",
+            "should_raise": False,
+        },
+    )
+    message = handler.start_pending("Flushing Pipe Network")
+    assert len(handler.messages) == 1
+    assert isinstance(handler.messages[0], PendingMessage)
+    assert not message.closed
+
+    handler.close_pending(message)
+    assert len(handler.messages) == 2
+    assert isinstance(handler.messages[1], ClosePendingMessage)
+    assert message.closed
+
+    # closing an already closed message does not lead to another broadcast
+    handler.close_pending(message)
+    assert len(handler.messages) == 2
+    assert message.closed
+
+
 # TODO
 # import json