Add a notebook test for get_document, fix rtc_add_doc_to_store for …

…notebooks

projects/jupyter-server-ydoc/jupyter_server_ydoc/pytest_plugin.py

@@ -149,11 +149,12 @@ async def _inner(format: str, type: str, path: str) -> Any:
 def rtc_add_doc_to_store(rtc_connect_doc_client):
     event = Event()
 
-    def _on_document_change(target: str, e: Any) -> None:
-        if target == "source":
-            event.set()
-
     async def _inner(format: str, type: str, path: str) -> None:
+        def _on_document_change(target: str, e: Any) -> None:
+            expected_target = "cells" if type == "notebook" else "source"
+            if target == expected_target:
+                event.set()
+
         if type == "notebook":
             doc = YNotebook()
         else:

tests/test_app.py

@@ -3,6 +3,9 @@
 
 from __future__ import annotations
 
+import json
+
+import nbformat
 import pytest
 from jupyter_server_ydoc.pytest_plugin import rtc_create_SQLite_store_factory
 from jupyter_server_ydoc.stores import SQLiteYStore, TempFileYStore
@@ -87,6 +90,25 @@ async def test_get_document_file(rtc_create_file, jp_serverapp, copy):
     await collaboration.stop_extension()
 
 
+@pytest.mark.parametrize("copy", [True, False])
+async def test_get_document_notebook(rtc_create_notebook, jp_serverapp, copy):
+    nb = nbformat.v4.new_notebook(
+        cells=[nbformat.v4.new_code_cell(source="1+1", execution_count=99)]
+    )
+    nb_content = nbformat.writes(nb, version=4)
+    path, _ = await rtc_create_notebook("test.ipynb", nb_content, store=True)
+    collaboration = jp_serverapp.web_app.settings["jupyter_server_ydoc"]
+    document = await collaboration.get_document(
+        path=path, content_type="notebook", file_format="json", copy=copy
+    )
+    doc = document.get()
+    assert len(doc["cells"]) == 1
+    cell = doc["cells"][0]
+    assert cell["source"] == "1+1"
+    assert cell["execution_count"] == 99
+    await collaboration.stop_extension()
+
+
 async def test_get_document_file_copy_is_independent(
     rtc_create_file, jp_serverapp, rtc_fetch_session
 ):