Add unit and integration tests for the kernels API

conftest.py

@@ -1,6 +1,12 @@
 import pytest
 
-pytest_plugins = ("pytest_jupyter.jupyter_server", "jupyter_server.pytest_plugin")
+pytest_plugins = ("pytest_jupyter.jupyter_server", "jupyter_server.pytest_plugin", "pytest_asyncio")
+
+
+def pytest_configure(config):
+    """Configure pytest settings."""
+    # Set asyncio fixture loop scope to function to avoid warnings
+    config.option.asyncio_default_fixture_loop_scope = "function"
 
 
 @pytest.fixture

jupyter_server_documents/kernels/kernel_client.py

@@ -80,7 +80,7 @@ async def stop_listening(self):
         """
         # If the listening task isn't defined yet
         # do nothing.
-        if not self._listening_task:
+        if not hasattr(self, '_listening_task') or not self._listening_task:
             return
 
         # Attempt to cancel the task.
@@ -93,6 +93,9 @@ async def stop_listening(self):
         # Log any exceptions that were raised.
         except Exception as err:
             self.log.error(err)
+        finally:
+            # Clear the task reference
+            self._listening_task = None
 
     _listening_task: t.Optional[t.Awaitable] = Any(allow_none=True)
 

jupyter_server_documents/kernels/kernel_manager.py

@@ -127,8 +127,9 @@ async def connect(self):
             await asyncio.sleep(0.1)
 
     async def disconnect(self):
-        await self.main_client.stop_listening()
-        self.main_client.stop_channels()
+        if self.main_client:
+            await self.main_client.stop_listening()
+            self.main_client.stop_channels()
 
     async def broadcast_state(self):
         """Broadcast state to all listeners"""

jupyter_server_documents/kernels/multi_kernel_manager.py

@@ -7,4 +7,11 @@ def start_watching_activity(self, kernel_id):
         pass
 
     def stop_buffering(self, kernel_id):
-        pass
+        pass
+
+    # NOTE: Since we disable watching activity and buffering here, 
+    # this method needs to be forked and remove code related to these things. 
+    async def restart_kernel(self, kernel_id, now=False):
+        """Restart a kernel by kernel_id"""
+        self._check_kernel_id(kernel_id)
+        await self.pinned_superclass._async_restart_kernel(self, kernel_id, now=now)

jupyter_server_documents/tests/kernels/conftest.py

@@ -0,0 +1,23 @@
+"""Configuration for kernel tests."""
+
+import pytest
+from unittest.mock import MagicMock
+
+
+@pytest.fixture
+def mock_logger():
+    """Create a mock logger for testing."""
+    return MagicMock()
+
+
+@pytest.fixture
+def mock_session():
+    """Create a mock session for testing."""
+    session = MagicMock()
+    session.msg_header.return_value = {"msg_id": "test-msg-id"}
+    session.msg.return_value = {"test": "message"}
+    session.serialize.return_value = ["", "serialized", "msg"]
+    session.deserialize.return_value = {"msg_type": "test", "content": b"test"}
+    session.unpack.return_value = {"test": "data"}
+    session.feed_identities.return_value = ([], [b"test", b"message"])
+    return session

jupyter_server_documents/tests/kernels/test_kernel_client.py

@@ -0,0 +1,105 @@
+import pytest
+from unittest.mock import MagicMock, patch
+
+from jupyter_server_documents.kernels.kernel_client import DocumentAwareKernelClient
+from jupyter_server_documents.kernels.message_cache import KernelMessageCache
+from jupyter_server_documents.outputs import OutputProcessor
+
+
+class TestDocumentAwareKernelClient:
+    """Test cases for DocumentAwareKernelClient."""
+
+    def test_default_message_cache(self):
+        """Test that message cache is created by default."""
+        client = DocumentAwareKernelClient()
+        assert isinstance(client.message_cache, KernelMessageCache)
+
+    def test_default_output_processor(self):
+        """Test that output processor is created by default."""
+        client = DocumentAwareKernelClient()
+        assert isinstance(client.output_processor, OutputProcessor)
+
+    @pytest.mark.asyncio
+    async def test_stop_listening_no_task(self):
+        """Test that stop_listening does nothing when no task exists."""
+        client = DocumentAwareKernelClient()
+        client._listening_task = None
+
+        # Should not raise an exception
+        await client.stop_listening()
+
+    def test_add_listener(self):
+        """Test adding a listener."""
+        client = DocumentAwareKernelClient()
+
+        def test_listener(channel, msg):
+            pass
+
+        client.add_listener(test_listener)
+
+        assert test_listener in client._listeners
+
+    def test_remove_listener(self):
+        """Test removing a listener."""
+        client = DocumentAwareKernelClient()
+
+        def test_listener(channel, msg):
+            pass
+
+        client.add_listener(test_listener)
+        client.remove_listener(test_listener)
+
+        assert test_listener not in client._listeners
+
+    @pytest.mark.asyncio
+    async def test_add_yroom(self):
+        """Test adding a YRoom."""
+        client = DocumentAwareKernelClient()
+
+        mock_yroom = MagicMock()
+        await client.add_yroom(mock_yroom)
+
+        assert mock_yroom in client._yrooms
+
+    @pytest.mark.asyncio
+    async def test_remove_yroom(self):
+        """Test removing a YRoom."""
+        client = DocumentAwareKernelClient()
+
+        mock_yroom = MagicMock()
+        client._yrooms.add(mock_yroom)
+
+        await client.remove_yroom(mock_yroom)
+
+        assert mock_yroom not in client._yrooms
+
+    def test_send_kernel_info_creates_message(self):
+        """Test that send_kernel_info creates a kernel info message."""
+        client = DocumentAwareKernelClient()
+
+        # Mock session
+        from jupyter_client.session import Session
+        client.session = Session()
+
+        with patch.object(client, 'handle_incoming_message') as mock_handle:
+            client.send_kernel_info()
+
+            # Verify that handle_incoming_message was called with shell channel
+            mock_handle.assert_called_once()
+            args, kwargs = mock_handle.call_args
+            assert args[0] == "shell"  # Channel name
+            assert isinstance(args[1], list)  # Message list
+
+    @pytest.mark.asyncio
+    async def test_handle_outgoing_message_control_channel(self):
+        """Test that control channel messages bypass document handling."""
+        client = DocumentAwareKernelClient()
+
+        msg = [b"test", b"message"]
+
+        with patch.object(client, 'handle_document_related_message') as mock_handle_doc:
+            with patch.object(client, 'send_message_to_listeners') as mock_send:
+                await client.handle_outgoing_message("control", msg)
+
+                mock_handle_doc.assert_not_called()
+                mock_send.assert_called_once_with("control", msg)