more tests

Author: craigbarratt

custom_components/pyscript/jupyter_kernel.py

@@ -196,6 +196,7 @@ def __init__(self, config, ast_ctx, global_ctx_name):
         self.ast_ctx = ast_ctx
 
         self.secure_key = str_to_bytes(self.config["key"])
+        self.no_connect_timeout = self.config.get("no_connect_timeout", 30)
         self.signature_schemes = {"hmac-sha256": hashlib.sha256}
         self.auth = hmac.HMAC(
             self.secure_key, digestmod=self.signature_schemes[self.config["signature_scheme"]],
@@ -297,21 +298,15 @@ async def send(
     ):
         """Send message to the Jupyter client."""
         header = self.new_header(msg_type)
-        if content is None:
-            content = {}
-        if parent_header is None:
-            parent_header = {}
-        if metadata is None:
-            metadata = {}
 
         def encode(msg):
             return str_to_bytes(json.dumps(msg))
 
         msg_lst = [
             encode(header),
-            encode(parent_header),
-            encode(metadata),
-            encode(content),
+            encode(parent_header if parent_header else {}),
+            encode(metadata if metadata else {}),
+            encode(content if content else {}),
         ]
         signature = self.msg_sign(msg_lst)
         parts = [DELIM, signature, msg_lst[0], msg_lst[1], msg_lst[2], msg_lst[3]]
@@ -447,6 +442,8 @@ async def shell_handler(self, shell_socket, wire_msg):
             )
 
         elif msg["header"]["msg_type"] == "complete_request":
+            root = ""
+            words = set()
             code = msg["content"]["code"]
             posn = msg["content"]["cursor_pos"]
             match = self.completion_re.match(code[0:posn].lower())
@@ -456,9 +453,6 @@ async def shell_handler(self, shell_socket, wire_msg):
                 words = words.union(await Function.service_completions(root))
                 words = words.union(await Function.func_completions(root))
                 words = words.union(self.ast_ctx.completions(root))
-            else:
-                root = ""
-                words = set()
             # _LOGGER.debug(f"complete_request code={code}, posn={posn}, root={root}, words={words}")
             content = {
                 "status": "ok",
@@ -708,8 +702,8 @@ async def housekeep_run(self):
     async def startup_timeout(self):
         """Shut down the session if nothing connects after 30 seconds."""
         await self.housekeep_q.put(["register", "startup_timeout", asyncio.current_task()])
-        await asyncio.sleep(30)
-        if self.task_cnt_max == 1:
+        await asyncio.sleep(self.no_connect_timeout)
+        if self.task_cnt_max <= 1:
             #
             # nothing started other than us, so shut down the session
             #

tests/test_jupyter.py

@@ -32,9 +32,9 @@ def msg_id():
     return str(uuid.uuid4())
 
 
-def msg_sign(msg_lst):
+def msg_sign(msg_lst, secret_key=SECRET_KEY):
     """Sign a message with a secure signature."""
-    auth_hmac = hmac.HMAC(SECRET_KEY, digestmod=hashlib.sha256)
+    auth_hmac = hmac.HMAC(secret_key, digestmod=hashlib.sha256)
     for msg in msg_lst:
         auth_hmac.update(msg)
     return str_to_bytes(auth_hmac.hexdigest())
@@ -73,27 +73,27 @@ def new_header(msg_type):
 
 
 async def send(
-    zmq_sock, msg_type, content=None, parent_header=None, metadata=None, identities=None,
+    zmq_sock,
+    msg_type,
+    content=None,
+    parent_header=None,
+    metadata=None,
+    identities=None,
+    secret_key=SECRET_KEY,
 ):
     """Send message to the Jupyter client."""
     header = new_header(msg_type)
-    if content is None:
-        content = {}
-    if parent_header is None:
-        parent_header = {}
-    if metadata is None:
-        metadata = {}
 
     def encode(msg):
         return str_to_bytes(json.dumps(msg))
 
     msg_lst = [
         encode(header),
-        encode(parent_header),
-        encode(metadata),
-        encode(content),
+        encode(parent_header if parent_header else {}),
+        encode(metadata if metadata else {}),
+        encode(content if content else {}),
     ]
-    signature = msg_sign(msg_lst)
+    signature = msg_sign(msg_lst, secret_key=secret_key)
     parts = [DELIM, signature, msg_lst[0], msg_lst[1], msg_lst[2], msg_lst[3]]
     if identities:
         parts = identities + parts
@@ -106,7 +106,7 @@ def encode(msg):
 PORT_NAMES = ["hb_port", "stdin_port", "shell_port", "iopub_port", "control_port"]
 
 
-async def setup_script(hass, now, source):
+async def setup_script(hass, now, source, no_connect=False):
     """Initialize and load the given pyscript."""
 
     scripts = [
@@ -149,6 +149,8 @@ def return_next_time():
         "signature_scheme": "hmac-sha256",
         "state_var": kernel_state_var,
     }
+    if no_connect:
+        kernel_cfg["no_connect_timeout"] = 0.0
     await hass.services.async_call("pyscript", "jupyter_kernel_start", kernel_cfg)
 
     while True:
@@ -160,6 +162,10 @@ def return_next_time():
     port_nums = json.loads(ports_state.state)
 
     sock = {}
+
+    if no_connect:
+        return sock, port_nums
+
     for name in PORT_NAMES:
         kernel_reader, kernel_writer = await asyncio.open_connection("127.0.0.1", port_nums[name])
         sock[name] = ZmqSocket(kernel_reader, kernel_writer, "ROUTER")
@@ -314,6 +320,10 @@ async def test_jupyter_kernel_msgs(hass, caplog):
     assert reply["header"]["msg_type"] == "is_complete_reply"
     assert reply["content"]["status"] == "invalid"
 
+    reply = await shell_msg(sock, "is_complete_request", {"code": "if 1:\n"})
+    assert reply["header"]["msg_type"] == "is_complete_reply"
+    assert reply["content"]["status"] == "incomplete"
+
     #
     # test code execution
     #
@@ -398,7 +408,23 @@ async def test_jupyter_kernel_port_close(hass, caplog):
         assert await sock["hb_port"].recv() == msg
 
     hass.bus.async_fire(EVENT_HOMEASSISTANT_STARTED)
-    await shutdown(sock)
+
+    #
+    # shut down the session via signature mismatch with bad key
+    #
+    await send(
+        sock["control_port"], "shutdown_request", {}, parent_header={}, identities={}, secret_key=b"bad_key"
+    )
+
+    #
+    # wait until the session ends, so the log receives the error message we check below
+    #
+    try:
+        await sock["iopub_port"].recv()
+    except EOFError:
+        pass
+
+    assert "signature mismatch: check_sig=" in caplog.text
 
 
 async def test_jupyter_kernel_redefine_func(hass, caplog):
@@ -474,3 +500,10 @@ async def test_jupyter_kernel_stdout(hass, caplog):
     assert stdout_msg["content"]["text"] == "hello\n"
 
     await shutdown(sock)
+
+
+async def test_jupyter_kernel_no_connection_timeout(hass, caplog):
+    """Test Jupyter kernel timeout on no connection."""
+    sock, port_nums = await setup_script(hass, [dt(2020, 7, 1, 11, 0, 0, 0)], "", no_connect=True)
+
+    assert "No connections to session " in caplog.text