diff --git a/crates/amalthea/src/fixtures/dummy_frontend.rs b/crates/amalthea/src/fixtures/dummy_frontend.rs index 9a303da19..32f82c3fe 100644 --- a/crates/amalthea/src/fixtures/dummy_frontend.rs +++ b/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 { diff --git a/crates/ark/src/fixtures/dummy_frontend.rs b/crates/ark/src/fixtures/dummy_frontend.rs index 07b9d3b51..172d0b354 100644 --- a/crates/ark/src/fixtures/dummy_frontend.rs +++ b/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> { - 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) + } +} diff --git a/crates/ark/tests/kernel-notebook.rs b/crates/ark/tests/kernel-notebook.rs new file mode 100644 index 000000000..08509be27 --- /dev/null +++ b/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); +} diff --git a/crates/ark/tests/kernel.rs b/crates/ark/tests/kernel.rs index fe9a93230..25d6a2bf8 100644 --- a/crates/ark/tests/kernel.rs +++ b/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); +}