Assert status field and return execution_count

Author: lionel-

crates/amalthea/src/fixtures/dummy_frontend.rs

@@ -160,21 +160,25 @@ impl DummyFrontend {
         Message::read_from_socket(&self.shell_socket).unwrap()
     }
 
-    /// Receive from Shell and assert ExecuteReply message
-    pub fn recv_shell_execute_reply(&self) -> Status {
+    /// Receive from Shell and assert `ExecuteReply` message.
+    /// Returns `execution_count`.
+    pub fn recv_shell_execute_reply(&self) -> u32 {
         let msg = self.recv_shell();
 
         assert_match!(msg, Message::ExecuteReply(data) => {
-            data.content.status
+            assert_eq!(data.content.status, Status::Ok);
+            data.content.execution_count
         })
     }
 
-    /// Receive from Shell and assert ExecuteReplyException message
-    pub fn recv_shell_execute_reply_exception(&self) -> Status {
-        let msg = Message::read_from_socket(&self.shell_socket).unwrap();
+    /// Receive from Shell and assert `ExecuteReplyException` message.
+    /// Returns `execution_count`.
+    pub fn recv_shell_execute_reply_exception(&self) -> u32 {
+        let msg = self.recv_shell();
 
         assert_match!(msg, Message::ExecuteReplyException(data) => {
-            data.content.status
+            assert_eq!(data.content.status, Status::Error);
+            data.content.execution_count
         })
     }
 

crates/ark/tests/kernel.rs

@@ -1,5 +1,4 @@
 use amalthea::wire::jupyter_message::Message;
-use amalthea::wire::jupyter_message::Status;
 use amalthea::wire::kernel_info_request::KernelInfoRequest;
 use ark::fixtures::DummyArkFrontend;
 use stdext::assert_match;
@@ -25,12 +24,13 @@ fn test_execute_request() {
     frontend.send_execute_request("42");
     frontend.recv_iopub_busy();
 
-    assert_eq!(frontend.recv_iopub_execute_input().code, "42");
+    let input = frontend.recv_iopub_execute_input();
+    assert_eq!(input.code, "42");
     assert_eq!(frontend.recv_iopub_execute_result(), "[1] 42");
 
     frontend.recv_iopub_idle();
 
-    assert_eq!(frontend.recv_shell_execute_reply(), Status::Ok);
+    assert_eq!(frontend.recv_shell_execute_reply(), input.execution_count)
 }
 
 #[test]
@@ -40,10 +40,14 @@ fn test_execute_request_error() {
     frontend.send_execute_request("stop('foobar')");
     frontend.recv_iopub_busy();
 
-    assert_eq!(frontend.recv_iopub_execute_input().code, "stop('foobar')");
+    let input = frontend.recv_iopub_execute_input();
+    assert_eq!(input.code, "stop('foobar')");
     assert!(frontend.recv_iopub_execute_error().contains("foobar"));
 
     frontend.recv_iopub_idle();
 
-    assert_eq!(frontend.recv_shell_execute_reply_exception(), Status::Error);
+    assert_eq!(
+        frontend.recv_shell_execute_reply_exception(),
+        input.execution_count
+    );
 }