diff --git a/conftest.py b/conftest.py index 6461459..ee29be3 100644 --- a/conftest.py +++ b/conftest.py @@ -1,12 +1,33 @@ import pytest -pytest_plugins = ("pytest_jupyter.jupyter_server",) +pytest_plugins = [ + "pytest_jupyter.jupyter_server", + "jupyter_server.pytest_plugin", + "jupyter_server_fileid.pytest_plugin", + "jupyter_server_ydoc.pytest_plugin" +] @pytest.fixture -def jp_server_config(jp_server_config): +def jp_server_config(jp_root_dir, jp_server_config): return { - "ServerApp": { - "jpserver_extensions": {"jupyter_server_nbmodel": True, "jupyter_server_ydoc": False} + 'ServerApp': { + 'jpserver_extensions': { + 'jupyter_server_ydoc': True, + 'jupyter_server_fileid': True, + 'jupyter_server_nbmodel': True, + }, + 'token': '', + 'password': '', + 'disable_check_xsrf': True + }, + "SQLiteYStore": {"db_path": str(jp_root_dir.joinpath(".rtc_test.db"))}, + "BaseFileIdManager": { + "root_dir": str(jp_root_dir), + "db_path": str(jp_root_dir.joinpath(".fid_test.db")), + "db_journal_mode": "OFF", + }, + "YDocExtension": { + "document_save_delay": 1 } } diff --git a/jupyter_server_nbmodel/extension.py b/jupyter_server_nbmodel/extension.py index ad945ff..34dd205 100644 --- a/jupyter_server_nbmodel/extension.py +++ b/jupyter_server_nbmodel/extension.py @@ -56,4 +56,5 @@ def initialize_handlers(self): async def stop_extension(self): if hasattr(self, "__execution_stack"): get_logger().info("Disposing the execution stackā€¦") + await self.__execution_stack.dispose() await asyncio.wait_for(self.__execution_stack.dispose(), timeout=3) diff --git a/jupyter_server_nbmodel/handlers.py b/jupyter_server_nbmodel/handlers.py index 2627f6e..3d36a88 100644 --- a/jupyter_server_nbmodel/handlers.py +++ b/jupyter_server_nbmodel/handlers.py @@ -8,6 +8,7 @@ from dataclasses import asdict, dataclass from functools import partial from http import HTTPStatus +from datetime import datetime, timezone import jupyter_server import jupyter_server.services @@ -123,7 +124,6 @@ async def _get_ycell( raise KeyError( msg, ) - return ycell @@ -219,6 +219,7 @@ async def _execute_snippet( The execution status and outputs. """ ycell = None + time_info = {} if metadata is not None: ycell = await _get_ycell(ydoc, metadata) if ycell is not None: @@ -227,6 +228,10 @@ async def _execute_snippet( del ycell["outputs"][:] ycell["execution_count"] = None ycell["execution_state"] = "running" + if metadata.get("record_timing", False): + time_info = ycell["metadata"].get("execution", {}) + time_info["shell.execute_reply.started"] = datetime.now(timezone.utc).isoformat()[:-6] + ycell["metadata"]["execution"] = time_info outputs = [] @@ -247,7 +252,13 @@ async def _execute_snippet( with ycell.doc.transaction(): ycell["execution_count"] = reply_content.get("execution_count") ycell["execution_state"] = "idle" - + if metadata and metadata.get("record_timing", False): + end_time = datetime.now(timezone.utc).isoformat()[:-6] + if reply_content["status"] == "ok": + time_info["shell.execute_reply"] = end_time + else: + time_info["execution_failed"] = end_time + ycell["metadata"]["execution"] = time_info return { "status": reply_content["status"], "execution_count": reply_content.get("execution_count"), @@ -524,9 +535,7 @@ async def post(self, kernel_id: str) -> None: msg = f"Unknown kernel with id: {kernel_id}" get_logger().error(msg) raise tornado.web.HTTPError(status_code=HTTPStatus.NOT_FOUND, reason=msg) - uid = self._execution_stack.put(kernel_id, snippet, metadata) - self.set_status(HTTPStatus.ACCEPTED) self.set_header("Location", f"/api/kernels/{kernel_id}/requests/{uid}") self.finish("{}") diff --git a/jupyter_server_nbmodel/tests/test_handlers.py b/jupyter_server_nbmodel/tests/test_handlers.py index eade4e2..b60d6ce 100644 --- a/jupyter_server_nbmodel/tests/test_handlers.py +++ b/jupyter_server_nbmodel/tests/test_handlers.py @@ -2,6 +2,7 @@ import datetime import json import re +import nbformat import pytest from jupyter_client.kernelspec import NATIVE_KERNEL_NAME @@ -149,6 +150,67 @@ async def test_post_erroneous_execute(jp_fetch, pending_kernel_is_ready, snippet await asyncio.sleep(1) +@pytest.mark.timeout(TEST_TIMEOUT) +async def test_execution_timing_metadata(jp_root_dir, jp_fetch, pending_kernel_is_ready, rtc_create_notebook, jp_serverapp): + snippet = "a = 1" + nb = nbformat.v4.new_notebook( + cells=[nbformat.v4.new_code_cell(source=snippet, execution_count=1)] + ) + 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"] + fim = jp_serverapp.web_app.settings["file_id_manager"] + document_id = f'json:notebook:{fim.get_id("test.ipynb")}' + cell_id = nb["cells"][0].get("id") + + r = await jp_fetch( + "api", "kernels", method="POST", body=json.dumps({"name": NATIVE_KERNEL_NAME}) + ) + kernel = json.loads(r.body.decode()) + await pending_kernel_is_ready(kernel["id"]) + + response = await wait_for_request( + jp_fetch, + "api", + "kernels", + kernel["id"], + "execute", + method="POST", + body=json.dumps({ + "code": snippet, + "metadata": { + "cell_id": cell_id, + "document_id": document_id, + "record_timing": True + } + }), + ) + assert response.code == 200 + + document = await collaboration.get_document( + path=path, content_type="notebook", file_format="json", copy=False + ) + cell_data = document.get()["cells"][0] + assert 'execution' in cell_data['metadata'], "'execution' does not exist in 'metadata'" + + # Assert that start and end time exist in 'execution' + execution = cell_data['metadata']['execution'] + assert 'shell.execute_reply.started' in execution, "'shell.execute_reply.started' does not exist in 'execution'" + assert 'shell.execute_reply' in execution, "'shell.execute_reply' does not exist in 'execution'" + + started_time = execution['shell.execute_reply.started'] + reply_time = execution['shell.execute_reply'] + + started_dt = datetime.datetime.fromisoformat(started_time) + reply_dt = datetime.datetime.fromisoformat(reply_time) + + # Assert that reply_time is greater than started_time + assert reply_dt > started_dt, "The reply time is not greater than the started time." + response2 = await jp_fetch("api", "kernels", kernel["id"], method="DELETE") + assert response2.code == 204 + await asyncio.sleep(1) + + @pytest.mark.timeout(TEST_TIMEOUT) async def test_post_input_execute(jp_fetch, pending_kernel_is_ready): # Start the first kernel diff --git a/pyproject.toml b/pyproject.toml index d49d5ce..d68f908 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,7 +30,14 @@ dynamic = ["version", "description", "authors", "urls", "keywords"] [project.optional-dependencies] lab = ["jupyterlab>=4.2.0", "jupyter-docprovider>=1.0.0b1", "jupyter-server-ydoc>=1.0.0b1"] -test = ["pytest~=8.2", "pytest-cov", "pytest-jupyter[server]>=0.6", "pytest-timeout"] +test = [ + "pytest~=8.2", + "pytest-cov", + "pytest-jupyter[server]>=0.6", + "pytest-timeout", + "jupyter-server-ydoc[test]>=1.0.0b1", + "jupyter-server-fileid" +] lint = ["mdformat>0.7", "mdformat-gfm>=0.3.5", "ruff>=0.4.0"] typing = ["mypy>=0.990"] @@ -74,10 +81,14 @@ build_dir = "jupyter_server_nbmodel/labextension" [tool.pytest.ini_options] filterwarnings = [ "error", + "ignore:Unclosed context