Add optional on_reset argument to YRoom get methods (#152)

* add on_reset callback to YRoom.get_{awareness,jupyter_ydoc,ydoc}

* add unit test for on_reset() callback

* fix mypy errors in yroom.py

Author: dlqqq

conftest.py

@@ -1,5 +1,17 @@
+from __future__ import annotations
 import pytest
 
+import asyncio
+import logging
+from traitlets.config import Config, LoggingConfigurable
+from jupyter_server.services.contents.filemanager import AsyncFileContentsManager
+from typing import TYPE_CHECKING
+from jupyter_server_documents.rooms.yroom_manager import YRoomManager
+
+if TYPE_CHECKING:
+    from jupyter_server.serverapp import ServerApp
+
+
 pytest_plugins = ("pytest_jupyter.jupyter_server", "jupyter_server.pytest_plugin", "pytest_asyncio")
 
 
@@ -10,5 +22,69 @@ def pytest_configure(config):
 
 
 @pytest.fixture
-def jp_server_config(jp_server_config):
-    return {"ServerApp": {"jpserver_extensions": {"jupyter_server_documents": True}}}
+def jp_server_config(jp_server_config, tmp_path):
+    """
+    Fixture that defines the traitlets configuration used in unit tests.
+    """
+
+    return Config({
+        "ServerApp": {
+            "jpserver_extensions": {
+                "jupyter_server_documents": True,
+                "jupyter_server_fileid": True
+            },
+            "root_dir": str(tmp_path)
+        },
+        "ContentsManager": {"root_dir": str(tmp_path)}
+    })
+
+class MockServerDocsApp(LoggingConfigurable):
+    """Mock `ServerDocsApp` class for testing purposes."""
+    
+    serverapp: ServerApp
+
+    def __init__(self, *args, serverapp: ServerApp, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.serverapp = serverapp
+        self._log = None
+        
+    @property
+    def log(self) -> logging.Logger:
+        return self.serverapp.log
+        
+    @property
+    def event_loop(self) -> asyncio.AbstractEventLoop:
+        return self.serverapp.io_loop.asyncio_loop
+    
+    @property
+    def contents_manager(self) -> AsyncFileContentsManager:
+        return self.serverapp.contents_manager
+
+
+@pytest.fixture
+def mock_server_docs_app(jp_server_config, jp_configurable_serverapp) -> MockServerDocsApp:
+    """
+    Returns a mocked `MockServerDocsApp` object that can be passed as the `parent`
+    argument to objects normally initialized by `ServerDocsApp` in `app.py`.
+    This should be passed to most of the "manager singletons" like
+    `YRoomManager`.
+
+    See `MockServerDocsApp` in `conftest.py` for a complete description of the
+    attributes, properties, and methods available. If something is missing,
+    please feel free to add to it in your PR.
+    
+    Returns:
+        A `MockServerDocsApp` instance that can be passed as the `parent` argument
+        to objects normally initialized by `ServerDocsApp`.
+    """
+    serverapp = jp_configurable_serverapp()
+    return MockServerDocsApp(config=jp_server_config, serverapp=serverapp)
+
+@pytest.fixture
+def mock_yroom_manager(mock_server_docs_app) -> YRoomManager:
+    """
+    Returns a mocked `YRoomManager` which can be passed as the `parent` argument
+    to `YRoom` for testing purposes.
+    """
+
+    return YRoomManager(parent=mock_server_docs_app)

jupyter_server_documents/rooms/yroom.py

@@ -144,6 +144,14 @@ class YRoom(LoggingConfigurable):
     `unobserve_jupyter_ydoc()`.
     """
 
+    # TODO: define a dataclass for this to ensure values are type-safe
+    _on_reset_callbacks: dict[Literal['awareness', 'ydoc', 'jupyter_ydoc'], list[Callable[[Any], Any]]]
+    """
+    Dictionary that stores all `on_reset` callbacks passed to `get_awareness()`,
+    `get_jupyter_ydoc()`, or `get_ydoc()`. These are stored in lists under the
+    'awareness', 'ydoc' and 'jupyter_ydoc' keys respectively.
+    """
+
     _ydoc: pycrdt.Doc
     """
     The `YDoc` for this room's document. See `get_ydoc()` documentation for more
@@ -197,6 +205,11 @@ def __init__(self, *args, **kwargs):
 
         # Initialize instance attributes
         self._jupyter_ydoc_observers = {}
+        self._on_reset_callbacks = {
+            "awareness": [],
+            "jupyter_ydoc": [],
+            "ydoc": [],
+        }
         self._stopped = False
         self._updated = False
         self._save_task = None
@@ -336,11 +349,16 @@ def clients(self) -> YjsClientGroup:
         return self._client_group
 
 
-    async def get_jupyter_ydoc(self) -> YBaseDoc:
+    async def get_jupyter_ydoc(self, on_reset: Callable[[YBaseDoc], Any] | None = None) -> YBaseDoc:
         """
-        Returns a reference to the room's JupyterYDoc
+        Returns a reference to the room's Jupyter YDoc
         (`jupyter_ydoc.ybasedoc.YBaseDoc`) after waiting for its content to be
         loaded from the ContentsManager.
+
+        This method also accepts an `on_reset` callback, which should take a
+        Jupyter YDoc as an argument. This callback is run with the new Jupyter
+        YDoc whenever the YDoc is reset, e.g. in response to an out-of-band
+        change.
         """
         if self.room_id == "JupyterLab:globalAwareness":
             message = "There is no Jupyter ydoc for global awareness scenario"
@@ -350,23 +368,39 @@ async def get_jupyter_ydoc(self) -> YBaseDoc:
             raise RuntimeError("Jupyter YDoc is not available")
         if self.file_api:
             await self.file_api.until_content_loaded
+        if on_reset:
+            self._on_reset_callbacks['jupyter_ydoc'].append(on_reset)
+            
         return self._jupyter_ydoc
 
 
-    async def get_ydoc(self) -> pycrdt.Doc:
+    async def get_ydoc(self, on_reset: Callable[[pycrdt.Doc], Any] | None = None) -> pycrdt.Doc:
         """
         Returns a reference to the room's YDoc (`pycrdt.Doc`) after
         waiting for its content to be loaded from the ContentsManager.
+
+        This method also accepts an `on_reset` callback, which should take a
+        YDoc as an argument. This callback is run with the new YDoc object
+        whenever the YDoc is reset, e.g. in response to an out-of-band change.
         """
         if self.file_api:
             await self.file_api.until_content_loaded
+        if on_reset:
+            self._on_reset_callbacks['ydoc'].append(on_reset)
         return self._ydoc
 
 
-    def get_awareness(self) -> pycrdt.Awareness:
+    def get_awareness(self, on_reset: Callable[[pycrdt.Awareness], Any] | None = None) -> pycrdt.Awareness:
         """
         Returns a reference to the room's awareness (`pycrdt.Awareness`).
+
+        This method also accepts an `on_reset` callback, which should take an
+        Awareness object as an argument. This callback is run with the new
+        Awareness object whenever the YDoc is reset, e.g. in response to an
+        out-of-band change.
         """
+        if on_reset:
+            self._on_reset_callbacks['awareness'].append(on_reset)
         return self._awareness
 
     def get_cell_execution_states(self) -> dict:
@@ -914,13 +948,34 @@ def _reset_ydoc(self) -> None:
         """
         Deletes and re-initializes the YDoc, awareness, and JupyterYDoc. This
         frees the memory occupied by their histories.
+
+        This runs all `on_reset` callbacks previously passed to `get_ydoc()`,
+        `get_jupyter_ydoc()`, or `get_awareness()`.
         """
         self._ydoc = self._init_ydoc()
         self._awareness = self._init_awareness(ydoc=self._ydoc)
         self._jupyter_ydoc = self._init_jupyter_ydoc(
             ydoc=self._ydoc,
             awareness=self._awareness
         )
+
+        # Run callbacks stored in `self._on_reset_callbacks`.
+        objects_by_type = {
+            "awareness": self._awareness,
+            "jupyter_ydoc": self._jupyter_ydoc,
+            "ydoc": self._ydoc,
+        }
+        for obj_type, obj in objects_by_type.items():
+            # This is type-safe, but requires a mypy hint because it cannot
+            # infer that `obj_type` only takes 3 values.
+            for on_reset in self._on_reset_callbacks[obj_type]: # type: ignore
+                try:
+                    result = on_reset(obj)
+                    if asyncio.iscoroutine(result):
+                        asyncio.create_task(result)
+                except Exception:
+                    self.log.exception(f"Exception raised by '{obj_type}' on_reset() callback:")
+                    continue
 
     @property
     def stopped(self) -> bool:

jupyter_server_documents/tests/test_yroom.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+import pytest
+import pytest_asyncio
+import os
+from unittest.mock import Mock
+from jupyter_server_documents.rooms.yroom import YRoom
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from pathlib import Path
+    from jupyter_server_documents.rooms import YRoomManager
+
+@pytest.fixture
+def mock_textfile_path(tmp_path: Path):
+    """
+    Returns the path of a mock text file under `/tmp`.
+
+    Automatically creates the file before each test & deletes the file after
+    each test.
+    """
+    # Create file before test and yield the path
+    path: Path = tmp_path / "test.txt"
+    path.touch()
+    yield path
+
+    # Cleanup after test
+    os.remove(path)
+
+
+@pytest_asyncio.fixture
+async def default_yroom(mock_yroom_manager: YRoomManager, mock_textfile_path: Path):
+    """
+    Returns a configured `YRoom` instance that serves an empty text file under
+    `/tmp`.
+
+    Uses the `mock_yroom_manager` fixture defined in `conftest.py`.
+    """
+    # Get room ID
+    file_id = mock_yroom_manager.fileid_manager.index(str(mock_textfile_path))
+    room_id = f"text:file:{file_id}"
+
+    # Initialize room and wait until its content is loaded
+    room: YRoom = YRoom(parent=mock_yroom_manager, room_id=room_id)
+    await room.file_api.until_content_loaded
+
+    # Yield configured `YRoom`
+    yield room
+
+    # Cleanup
+    room.stop(immediately=True)
+
+class TestDefaultYRoom():
+    """
+    Tests that assert against the `default_yroom` fixture defined above.
+    """
+
+    @pytest.mark.asyncio
+    async def test_on_reset_callbacks(self, default_yroom: YRoom):
+        """
+        Asserts that the `on_reset()` callback passed to
+        `YRoom.get_{awareness,jupyter_ydoc,ydoc}()` methods are each called with
+        the expected object when the YDoc is reset.
+        """
+        yroom = default_yroom
+        
+        # Create mock callbacks
+        awareness_reset_mock = Mock()
+        jupyter_ydoc_reset_mock = Mock()
+        ydoc_reset_mock = Mock()
+        
+        # Call get methods while passing `on_reset` callbacks
+        yroom.get_awareness(on_reset=awareness_reset_mock)
+        await yroom.get_jupyter_ydoc(on_reset=jupyter_ydoc_reset_mock)
+        await yroom.get_ydoc(on_reset=ydoc_reset_mock)
+        
+        # Assert that each callback has not been called yet
+        awareness_reset_mock.assert_not_called()
+        jupyter_ydoc_reset_mock.assert_not_called()
+        ydoc_reset_mock.assert_not_called()
+        
+        # Reset the ydoc and get the new expected objects
+        yroom._reset_ydoc()
+        new_awareness = yroom.get_awareness()
+        new_jupyter_ydoc = await yroom.get_jupyter_ydoc()
+        new_ydoc = await yroom.get_ydoc()
+        
+        # Assert that each callback was called exactly once with the expected
+        # object
+        awareness_reset_mock.assert_called_once_with(new_awareness)
+        jupyter_ydoc_reset_mock.assert_called_once_with(new_jupyter_ydoc)
+        ydoc_reset_mock.assert_called_once_with(new_ydoc)
+

jupyter_server_documents/tests/test_yroom_file_api.py

@@ -54,8 +54,8 @@ def mock_plaintext_file(tmp_path):
     os.remove(target_path)
 
 
-@pytest_asyncio.fixture(loop_scope="module")
-async def plaintext_file_api(
+@pytest.fixture
+def plaintext_file_api(
     mock_plaintext_file: str,
     jp_contents_manager: AsyncFileContentsManager,
     fileid_manager: ArbitraryFileIdManager

pyproject.toml

@@ -39,7 +39,9 @@ dynamic = ["version", "description", "authors", "urls", "keywords"]
 test = [
     "coverage",
     "pytest",
-    "pytest-asyncio",
+    # pytest-asyncio <1.1.0 incorrectly resolves async fixtures to async
+    # generators instead of the yielded object.
+    "pytest-asyncio>=1.1.0",
     "pytest-cov",
     "pytest-jupyter[server]>=0.6.0",
 ]