Fix unit tests

Author: fcollonval

jupyter_server_nbmodel/handlers.py

@@ -82,7 +82,9 @@ def get(self, kernel_id: str, uid: str) -> t.Any:
         else:
             return None
 
-    def put(self, km: jupyter_client.manager.KernelManager, snippet: str, ycell: y.Map) -> str:
+    def put(
+        self, km: jupyter_client.manager.KernelManager, snippet: str, ycell: y.Map | None
+    ) -> str:
         """Add a asynchronous execution request.
 
         Args:
@@ -99,18 +101,27 @@ def put(self, km: jupyter_client.manager.KernelManager, snippet: str, ycell: y.M
         )
         return uid
 
-    def _stdin_hook(self, kernel_id, msg) -> None:
+    def _stdin_hook(self, kernel_id: str, msg: dict) -> None:
         get_logger().info(f"Execution request {kernel_id} received a input request {msg!s}")
         if kernel_id in self.__pending_inputs:
-            get_logger().error(f"Execution request {kernel_id} received a input request while waiting for an input.\n{msg}")
+            get_logger().error(
+                f"Execution request {kernel_id} received a input request while waiting for an input.\n{msg}"  # noqa: E501
+            )
 
         header = msg["header"].copy()
         header["date"] = header["date"].isoformat()
-        self.__pending_inputs[kernel_id] = {"parent_header": header, "input_request": msg["content"]}
+        self.__pending_inputs[kernel_id] = {
+            "parent_header": header,
+            "input_request": msg["content"],
+        }
 
 
 async def execute_task(
-    uid, km: jupyter_client.manager.KernelManager, snippet: str, ycell: y.Map, stdin_hook
+    uid,
+    km: jupyter_client.manager.KernelManager,
+    snippet: str,
+    ycell: y.Map | None,
+    stdin_hook: t.Callable[[dict], None] | None,
 ) -> t.Any:
     try:
         get_logger().debug(f"Will execute request {uid}.")
@@ -134,15 +145,11 @@ async def execute_task(
 
 async def _execute_snippet(
     uid: str,
-    km: jupyter_client.client.KernelClient,
+    km: jupyter_client.manager.KernelManager,
     snippet: str,
-    ycell: y.Map,
-    stdin_hook,
+    ycell: y.Map | None,
+    stdin_hook: t.Callable[[dict], None] | None,
 ) -> dict[str, t.Any]:
-    client = km.client()
-    client.session.session = uid
-    # FIXME
-    # client.session.username = username
 
     if ycell is not None:
         # Reset cell
@@ -151,6 +158,10 @@ async def _execute_snippet(
             ycell["execution_count"] = None
 
     outputs = []
+    client = km.client()
+    client.session.session = uid
+    # FIXME
+    # client.session.username = username
 
     # FIXME we don't check if the session is consistent (aka the kernel is linked to the document)
     #   - should we?
@@ -175,6 +186,7 @@ async def _execute_snippet(
             "outputs": json.dumps(outputs),
         }
     finally:
+        client.stop_channels()
         del client
 
 
@@ -252,6 +264,7 @@ async def post(self, kernel_id: str) -> None:
         body = self.get_json_body()
 
         snippet = body.get("code")
+        ycell = None
         # From RTC model
         if snippet is None:
             document_id = body.get("document_id")
@@ -280,7 +293,6 @@ async def post(self, kernel_id: str) -> None:
                 raise tornado.web.HTTPError(status_code=HTTPStatus.NOT_FOUND, reason=msg)
 
             ycells = filter(lambda c: c["id"] == cell_id, notebook.ycells)
-            ycell = None
             try:
                 ycell = next(ycells)
             except StopIteration:
@@ -337,9 +349,12 @@ async def post(self, kernel_id: str) -> None:
         try:
             # only send stdin reply if there *was not* another request
             # or execution finished while we were reading.
-            if not (await client.stdin_channel.msg_ready() or await client.shell_channel.msg_ready()):
+            if not (
+                await client.stdin_channel.msg_ready() or await client.shell_channel.msg_ready()
+            ):
                 client.input(body["input"])
         finally:
+            client.stop_channels()
             del client
 
 

jupyter_server_nbmodel/tests/test_handlers.py

@@ -1,11 +1,41 @@
 import asyncio
+import datetime
 import json
-
-from jupyter_client.kernelspec import NATIVE_KERNEL_NAME
+import re
 
 import pytest
+from jupyter_client.kernelspec import NATIVE_KERNEL_NAME
 
 TEST_TIMEOUT = 60
+SLEEP = 0.1
+
+
+REQUEST_REGEX = re.compile(r"^/api/kernels/\w+-\w+-\w+-\w+-\w+/requests/\w+-\w+-\w+-\w+-\w+$")
+
+
+async def _wait_request(fetch, endpoint: str):
+    """Poll periodically to fetch the execution request result."""
+    start_time = datetime.datetime.now()
+
+    while (datetime.datetime.now() - start_time).total_seconds() < 0.9 * TEST_TIMEOUT:
+        await asyncio.sleep(SLEEP)
+        response = await fetch(endpoint)
+        response.rethrow()
+        if response.code != 202:
+            return response
+
+    raise TimeoutError(f"Request {endpoint} timed out.")
+
+
+async def wait_for_request(fetch, *args, **kwargs):
+    """Wait for execution request."""
+    r = await fetch(*args, **kwargs)
+    assert r.code == 202
+    location = r.headers["Location"]
+    assert REQUEST_REGEX.match(location) is not None
+
+    ans = await _wait_request(fetch, location)
+    return ans
 
 
 @pytest.fixture()
@@ -39,7 +69,7 @@ async def _(kernel_id, ready=None):
         (
             """from IPython.display import HTML
 HTML('<p><b>Jupyter</b> rocks.</p>')""",
-            '{"output_type": "execute_result", "metadata": {}, "data": {"text/plain": "<IPython.core.display.HTML object>", "text/html": "<p><b>Jupyter</b> rocks.</p>"}, "execution_count": 1}',
+            '{"output_type": "execute_result", "metadata": {}, "data": {"text/plain": "<IPython.core.display.HTML object>", "text/html": "<p><b>Jupyter</b> rocks.</p>"}, "execution_count": 1}',  # noqa: E501
         ),
     ),
 )
@@ -50,7 +80,8 @@ async def test_post_execute(jp_fetch, pending_kernel_is_ready, snippet, output):
     kernel = json.loads(r.body.decode())
     await pending_kernel_is_ready(kernel["id"])
 
-    response = await jp_fetch(
+    response = await wait_for_request(
+        jp_fetch,
         "api",
         "kernels",
         kernel["id"],
@@ -89,7 +120,8 @@ async def test_post_erroneous_execute(jp_fetch, pending_kernel_is_ready, snippet
     kernel = json.loads(r.body.decode())
     await pending_kernel_is_ready(kernel["id"])
 
-    response = await jp_fetch(
+    response = await wait_for_request(
+        jp_fetch,
         "api",
         "kernels",
         kernel["id"],