Test multiple expressions in console and notebook modes (#559)

crates/amalthea/src/fixtures/dummy_frontend.rs

@@ -18,6 +18,7 @@ use crate::wire::jupyter_message::Message;
 use crate::wire::jupyter_message::ProtocolMessage;
 use crate::wire::jupyter_message::Status;
 use crate::wire::status::ExecutionState;
+use crate::wire::stream::Stream;
 use crate::wire::wire_message::WireMessage;
 
 pub struct DummyFrontend {
@@ -247,6 +248,24 @@ impl DummyFrontend {
         })
     }
 
+    pub fn recv_iopub_stream_stdout(&self) -> String {
+        let msg = self.recv_iopub();
+
+        assert_matches!(msg, Message::StreamOutput(data) => {
+            assert_eq!(data.content.name, Stream::Stdout);
+            data.content.text
+        })
+    }
+
+    pub fn recv_iopub_stream_stderr(&self) -> String {
+        let msg = self.recv_iopub();
+
+        assert_matches!(msg, Message::StreamOutput(data) => {
+            assert_eq!(data.content.name, Stream::Stderr);
+            data.content.text
+        })
+    }
+
     /// Receive from IOPub and assert ExecuteResult message. Returns compulsory
     /// `evalue` field.
     pub fn recv_iopub_execute_error(&self) -> String {

crates/ark/src/fixtures/dummy_frontend.rs

@@ -23,6 +23,16 @@ pub struct DummyArkFrontend {
     guard: MutexGuard<'static, DummyFrontend>,
 }
 
+/// Wrapper around `DummyArkFrontend` that uses `SessionMode::Notebook`
+///
+/// Only one of `DummyArkFrontend` or `DummyArkFrontendNotebook` can be used in
+/// a given process. Just don't import both and you should be fine as Rust will
+/// let you know about a missing symbol if you happen to copy paste `lock()`
+/// calls of different kernel types between files.
+pub struct DummyArkFrontendNotebook {
+    inner: DummyArkFrontend,
+}
+
 impl DummyArkFrontend {
     pub fn lock() -> Self {
         Self {
@@ -31,10 +41,13 @@ impl DummyArkFrontend {
     }
 
     fn get_frontend() -> &'static Arc<Mutex<DummyFrontend>> {
-        FRONTEND.get_or_init(|| Arc::new(Mutex::new(DummyArkFrontend::init())))
+        // These are the hard-coded defaults. Call `init()` explicitly to
+        // override.
+        let session_mode = SessionMode::Console;
+        FRONTEND.get_or_init(|| Arc::new(Mutex::new(DummyArkFrontend::init(session_mode))))
     }
 
-    fn init() -> DummyFrontend {
+    pub(crate) fn init(session_mode: SessionMode) -> DummyFrontend {
         if FRONTEND.get().is_some() {
             panic!("Can't spawn Ark more than once");
         }
@@ -52,7 +65,7 @@ impl DummyArkFrontend {
                 String::from("--no-restore"),
             ],
             None,
-            SessionMode::Console,
+            session_mode,
             false,
         );
 
@@ -87,3 +100,38 @@ impl DerefMut for DummyArkFrontend {
         DerefMut::deref_mut(&mut self.guard)
     }
 }
+
+impl DummyArkFrontendNotebook {
+    /// Lock a notebook frontend.
+    ///
+    /// NOTE: Only one of `DummyArkFrontendNotebook::lock()` re
+    /// `DummyArkFrontend::lock()` should be called in a given process.
+    pub fn lock() -> Self {
+        Self::init();
+
+        Self {
+            inner: DummyArkFrontend::lock(),
+        }
+    }
+
+    /// Initialize with Notebook session mode
+    fn init() {
+        let session_mode = SessionMode::Notebook;
+        FRONTEND.get_or_init(|| Arc::new(Mutex::new(DummyArkFrontend::init(session_mode))));
+    }
+}
+
+// Allow method calls to be forwarded to inner type
+impl Deref for DummyArkFrontendNotebook {
+    type Target = DummyFrontend;
+
+    fn deref(&self) -> &Self::Target {
+        Deref::deref(&self.inner)
+    }
+}
+
+impl DerefMut for DummyArkFrontendNotebook {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        DerefMut::deref_mut(&mut self.inner)
+    }
+}

crates/ark/tests/kernel-notebook.rs

@@ -0,0 +1,39 @@
+use ark::fixtures::DummyArkFrontendNotebook;
+
+#[test]
+fn test_notebook_execute_request() {
+    let frontend = DummyArkFrontendNotebook::lock();
+
+    frontend.send_execute_request("42");
+    frontend.recv_iopub_busy();
+
+    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(), input.execution_count);
+}
+
+#[test]
+fn test_notebook_execute_request_multiple_expressions() {
+    let frontend = DummyArkFrontendNotebook::lock();
+
+    let code = "1\nprint(2)\n3";
+    frontend.send_execute_request(code);
+    frontend.recv_iopub_busy();
+
+    let input = frontend.recv_iopub_execute_input();
+    assert_eq!(input.code, code);
+
+    // Printed output
+    assert_eq!(frontend.recv_iopub_stream_stdout(), "[1] 2\n");
+
+    // Unlike console mode, we don't get intermediate results in notebooks
+    assert_eq!(frontend.recv_iopub_execute_result(), "[1] 3");
+
+    frontend.recv_iopub_idle();
+
+    assert_eq!(frontend.recv_shell_execute_reply(), input.execution_count);
+}

crates/ark/tests/kernel.rs

@@ -54,3 +54,27 @@ fn test_execute_request_error() {
         input.execution_count
     );
 }
+
+#[test]
+fn test_execute_request_multiple_expressions() {
+    let frontend = DummyArkFrontend::lock();
+
+    let code = "1\nprint(2)\n3";
+    frontend.send_execute_request(code);
+    frontend.recv_iopub_busy();
+
+    let input = frontend.recv_iopub_execute_input();
+    assert_eq!(input.code, code);
+
+    // Printed output
+    assert_eq!(frontend.recv_iopub_stream_stdout(), "[1] 2\n");
+
+    // In console mode, we get output for all intermediate results.  That's not
+    // the case in notebook mode where only the final result is emitted. Note
+    // that `print()` returns invisibly.
+    assert_eq!(frontend.recv_iopub_execute_result(), "[1] 1\n[1] 3");
+
+    frontend.recv_iopub_idle();
+
+    assert_eq!(frontend.recv_shell_execute_reply(), input.execution_count);
+}