Merge pull request #5 from runtimed/comms-testing

Implement comms_lifecycle test

Author: rgbkrk

src/harness.rs

@@ -5,8 +5,9 @@ use crate::types::{KernelReport, TestCategory, TestRecord, TestResult};
 use chrono::Utc;
 use jupyter_protocol::connection_info::{ConnectionInfo, Transport};
 use jupyter_protocol::messaging::{
-    ExecuteRequest, ExecutionState, InputReply, JupyterMessage, JupyterMessageContent,
-    KernelInfoReply, KernelInfoRequest, ReplyStatus, ShutdownRequest, Status,
+    CommClose, CommOpen, ExecuteRequest, ExecutionState, InputReply, JupyterMessage,
+    JupyterMessageContent, KernelInfoReply, KernelInfoRequest, ReplyStatus, ShutdownRequest,
+    Status,
 };
 use runtimelib::{
     create_client_control_connection, create_client_heartbeat_connection,
@@ -224,6 +225,60 @@ impl KernelUnderTest {
             .map_err(|e| HarnessError::ProtocolError(e.to_string()))
     }
 
+    /// Send a request on shell and wait for reply, also collecting IOPub messages.
+    pub async fn shell_request_with_iopub(
+        &mut self,
+        content: impl Into<JupyterMessageContent>,
+    ) -> Result<(JupyterMessage, Vec<JupyterMessage>)> {
+        let request: JupyterMessage = JupyterMessage::new(content, None);
+        let msg_id = request.header.msg_id.clone();
+
+        self.shell
+            .send(request)
+            .await
+            .map_err(|e| HarnessError::ProtocolError(e.to_string()))?;
+
+        // Collect IOPub messages until idle
+        let mut iopub_messages = Vec::new();
+        let start = Instant::now();
+
+        loop {
+            if start.elapsed() > self.test_timeout {
+                return Err(HarnessError::Timeout("iopub idle".to_string()));
+            }
+
+            match timeout(Duration::from_millis(100), self.iopub.read()).await {
+                Ok(Ok(msg)) => {
+                    if msg.parent_header.as_ref().map(|h| &h.msg_id) == Some(&msg_id) {
+                        let is_idle = matches!(
+                            &msg.content,
+                            JupyterMessageContent::Status(Status { execution_state })
+                            if *execution_state == ExecutionState::Idle
+                        );
+                        iopub_messages.push(msg);
+                        if is_idle {
+                            break;
+                        }
+                    }
+                }
+                Ok(Err(e)) => {
+                    return Err(HarnessError::ProtocolError(e.to_string()));
+                }
+                Err(_) => {
+                    // Timeout on this read, continue
+                }
+            }
+        }
+
+        // Read shell reply
+        let reply = timeout(self.test_timeout, self.shell.read())
+            .await
+            .map_err(|_| HarnessError::Timeout("shell reply".to_string()))?
+            .map_err(|e| HarnessError::ProtocolError(e.to_string()))?;
+
+        Ok((reply, iopub_messages))
+    }
+
     /// Send a request on control and wait for reply.
     pub async fn control_request(
         &mut self,
@@ -398,6 +453,47 @@ impl KernelUnderTest {
         &mut self.stdin
     }
 
+    /// Send comm_open and check if kernel rejects it (returns true if rejected).
+    pub async fn send_comm_open(&mut self, msg: CommOpen) -> Result<bool> {
+        let comm_id = msg.comm_id.clone();
+        let request: JupyterMessage = JupyterMessage::new(msg, None);
+
+        self.shell
+            .send(request)
+            .await
+            .map_err(|e| HarnessError::ProtocolError(e.to_string()))?;
+
+        // Brief wait for potential comm_close rejection on IOPub
+        let start = Instant::now();
+        while start.elapsed() < Duration::from_millis(500) {
+            match timeout(Duration::from_millis(100), self.iopub.read()).await {
+                Ok(Ok(msg)) => {
+                    if let JupyterMessageContent::CommClose(close) = &msg.content {
+                        if close.comm_id == comm_id {
+                            return Ok(true); // Rejected
+                        }
+                    }
+                }
+                _ => {}
+            }
+        }
+
+        Ok(false) // Not rejected (accepted or ignored)
+    }
+
+    /// Send comm_close to clean up a comm.
+    pub async fn send_comm_close(&mut self, msg: CommClose) -> Result<()> {
+        let request: JupyterMessage = JupyterMessage::new(msg, None);
+        self.shell
+            .send(request)
+            .await
+            .map_err(|e| HarnessError::ProtocolError(e.to_string()))?;
+
+        // Brief wait for processing
+        tokio::time::sleep(Duration::from_millis(100)).await;
+        Ok(())
+    }
+
     /// Shutdown the kernel cleanly.
     pub async fn shutdown(mut self) -> Result<()> {
         let request = ShutdownRequest { restart: false };

src/tests.rs

@@ -3,9 +3,9 @@
 use crate::harness::{ConformanceTest, KernelUnderTest};
 use crate::types::{FailureKind, TestCategory, TestResult};
 use jupyter_protocol::messaging::{
-    CommInfoRequest, CompleteRequest, ExecutionState, HistoryRequest, InspectRequest,
-    InterruptRequest, IsCompleteReplyStatus, IsCompleteRequest, JupyterMessageContent, ReplyStatus,
-    ShutdownRequest, Status, StreamContent,
+    CommClose, CommId, CommInfoRequest, CommOpen, CompleteRequest, ExecutionState, HistoryRequest,
+    InspectRequest, InterruptRequest, IsCompleteReplyStatus, IsCompleteRequest,
+    JupyterMessageContent, ReplyStatus, ShutdownRequest, Status, StreamContent,
 };
 use std::future::Future;
 use std::pin::Pin;
@@ -618,9 +618,53 @@ fn test_stdin_input_request(
 }
 
 fn test_comms_lifecycle(
-    _kernel: &mut KernelUnderTest,
+    kernel: &mut KernelUnderTest,
 ) -> Pin<Box<dyn Future<Output = TestResult> + Send + '_>> {
-    Box::pin(async move { TestResult::Unsupported })
+    Box::pin(async move {
+        // Generate a unique comm_id for this test
+        let comm_id = CommId(format!("test-comm-{}", uuid::Uuid::new_v4()));
+
+        // Open a comm with a test target
+        // Most kernels won't have this registered, but should handle it gracefully
+        let open_msg = CommOpen {
+            comm_id: comm_id.clone(),
+            target_name: "jupyter.kernel_testbed.test".to_string(),
+            data: serde_json::Map::new(),
+            target_module: None,
+        };
+
+        // Send comm_open on shell channel
+        // Note: comm_open doesn't get a direct reply, but we can verify the kernel
+        // handles it by checking IOPub for comm_close (rejection) or doing a
+        // follow-up execute to ensure kernel is still responsive
+        match kernel.send_comm_open(open_msg).await {
+            Ok(rejection) => {
+                if rejection {
+                    // Kernel properly rejected unknown target with comm_close
+                    TestResult::Pass
+                } else {
+                    // Kernel accepted (or silently ignored) the comm_open
+                    // Send comm_close to clean up, then verify kernel still works
+                    let close_msg = CommClose {
+                        comm_id: comm_id.clone(),
+                        data: serde_json::Map::new(),
+                    };
+                    let _ = kernel.send_comm_close(close_msg).await;
+
+                    // Verify kernel is still responsive
+                    let code = kernel.snippets().complete_code.to_string();
+                    match kernel.execute_and_collect(&code).await {
+                        Ok(_) => TestResult::Pass,
+                        Err(e) => TestResult::fail(
+                            format!("Kernel unresponsive after comm: {}", e),
+                            FailureKind::HarnessError,
+                        ),
+                    }
+                }
+            }
+            Err(e) => TestResult::fail(e.to_string(), FailureKind::HarnessError),
+        }
+    })
 }
 
 fn test_interrupt_request(