Add test suite and minor refinements to the utility subsystem

node/subsystem/Cargo.toml

@@ -20,3 +20,9 @@ polkadot-statement-table = { path = "../../statement-table" }
 sc-network = { git = "https://github.com/paritytech/substrate", branch = "master" }
 sp-core = { git = "https://github.com/paritytech/substrate", branch = "master" }
 streamunordered = "0.5.1"
+
+[dev-dependencies]
+assert_matches = "1.3.0"
+async-trait = "0.1"
+futures = { version = "0.3.5", features = ["thread-pool"] }
+parking_lot = "0.10.0"

node/subsystem/src/lib.rs

@@ -35,6 +35,8 @@ use crate::messages::AllMessages;
 
 pub mod messages;
 pub mod util;
+#[cfg(test)]
+pub mod test_helpers;
 
 /// Signals sent by an overseer to a subsystem.
 #[derive(PartialEq, Clone, Debug)]

node/subsystem/src/messages.rs

@@ -408,4 +408,10 @@ pub enum AllMessages {
 	AvailabilityStore(AvailabilityStoreMessage),
 	/// Message for the network bridge subsystem.
 	NetworkBridge(NetworkBridgeMessage),
+	/// Test message
+	///
+	/// This variant is only valid while testing, but makes the process of testing the
+	/// subsystem job manager much simpler.
+	#[cfg(test)]
+	Test(String),
 }

node/subsystem/src/test_helpers.rs

@@ -0,0 +1,229 @@
+// Copyright 2017-2020 Parity Technologies (UK) Ltd.
+// This file is part of Polkadot.
+
+// Polkadot is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Polkadot is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Polkadot.  If not, see <http://www.gnu.org/licenses/>.
+
+//! Utilities for testing subsystems.
+
+use crate::{SubsystemContext, FromOverseer, SubsystemResult, SubsystemError};
+use crate::messages::AllMessages;
+
+use futures::prelude::*;
+use futures::channel::mpsc;
+use futures::task::{Spawn, SpawnExt};
+use futures::poll;
+use parking_lot::Mutex;
+
+use std::convert::Infallible;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll, Waker};
+
+enum SinkState<T> {
+	Empty {
+		read_waker: Option<Waker>,
+	},
+	Item {
+		item: T,
+		ready_waker: Option<Waker>,
+		flush_waker: Option<Waker>,
+	},
+}
+
+/// The sink half of a single-item sink that does not resolve until the item has been read.
+pub struct SingleItemSink<T>(Arc<Mutex<SinkState<T>>>);
+
+/// The stream half of a single-item sink.
+pub struct SingleItemStream<T>(Arc<Mutex<SinkState<T>>>);
+
+impl<T> Sink<T> for SingleItemSink<T> {
+	type Error = Infallible;
+
+	fn poll_ready(
+		self: Pin<&mut Self>,
+		cx: &mut Context,
+	) -> Poll<Result<(), Infallible>> {
+		let mut state = self.0.lock();
+		match *state {
+			SinkState::Empty { .. } => Poll::Ready(Ok(())),
+			SinkState::Item { ref mut ready_waker, .. } => {
+				*ready_waker = Some(cx.waker().clone());
+				Poll::Pending
+			}
+		}
+	}
+
+	fn start_send(
+		self: Pin<&mut Self>,
+		item: T,
+	) -> Result<(), Infallible> {
+		let mut state = self.0.lock();
+
+		match *state {
+			SinkState::Empty { ref mut read_waker } => {
+				if let Some(waker) = read_waker.take() {
+					waker.wake();
+				}
+			}
+			_ => panic!("start_send called outside of empty sink state ensured by poll_ready"),
+		}
+
+		*state = SinkState::Item {
+			item,
+			ready_waker: None,
+			flush_waker: None,
+		};
+
+		Ok(())
+	}
+
+	fn poll_flush(
+		self: Pin<&mut Self>,
+		cx: &mut Context,
+	) -> Poll<Result<(), Infallible>> {
+		let mut state = self.0.lock();
+		match *state {
+			SinkState::Empty { .. } => Poll::Ready(Ok(())),
+			SinkState::Item { ref mut flush_waker, .. } => {
+				*flush_waker = Some(cx.waker().clone());
+				Poll::Pending
+			}
+		}
+	}
+
+	fn poll_close(
+		self: Pin<&mut Self>,
+		cx: &mut Context,
+	) -> Poll<Result<(), Infallible>> {
+		self.poll_flush(cx)
+	}
+}
+
+impl<T> Stream for SingleItemStream<T> {
+	type Item = T;
+
+	fn poll_next(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
+		let mut state = self.0.lock();
+
+		let read_waker = Some(cx.waker().clone());
+
+		match std::mem::replace(&mut *state, SinkState::Empty { read_waker }) {
+			SinkState::Empty { .. } => Poll::Pending,
+			SinkState::Item { item, ready_waker, flush_waker } => {
+				if let Some(waker) = ready_waker {
+					waker.wake();
+				}
+
+				if let Some(waker) = flush_waker {
+					waker.wake();
+				}
+
+				Poll::Ready(Some(item))
+			}
+		}
+	}
+}
+
+/// Create a single-item Sink/Stream pair.
+///
+/// The sink's send methods resolve at the point which the stream reads the item,
+/// not when the item is buffered.
+pub fn single_item_sink<T>() -> (SingleItemSink<T>, SingleItemStream<T>) {
+	let inner = Arc::new(Mutex::new(SinkState::Empty { read_waker: None }));
+	(
+		SingleItemSink(inner.clone()),
+		SingleItemStream(inner),
+	)
+}
+
+/// A test subsystem context.
+pub struct TestSubsystemContext<M, S> {
+	tx: mpsc::UnboundedSender<AllMessages>,
+	rx: SingleItemStream<FromOverseer<M>>,
+	spawn: S,
+}
+
+#[async_trait::async_trait]
+impl<M: Send + 'static, S: Spawn + Send + 'static> SubsystemContext for TestSubsystemContext<M, S> {
+	type Message = M;
+
+	async fn try_recv(&mut self) -> Result<Option<FromOverseer<M>>, ()> {
+		match poll!(self.rx.next()) {
+			Poll::Ready(Some(msg)) => Ok(Some(msg)),
+			Poll::Ready(None) => Err(()),
+			Poll::Pending => Ok(None),
+		}
+	}
+
+	async fn recv(&mut self) -> SubsystemResult<FromOverseer<M>> {
+		self.rx.next().await.ok_or(SubsystemError)
+	}
+
+	async fn spawn(&mut self, s: Pin<Box<dyn Future<Output = ()> + Send>>) -> SubsystemResult<()> {
+		self.spawn.spawn(s).map_err(Into::into)
+	}
+
+	async fn send_message(&mut self, msg: AllMessages) -> SubsystemResult<()> {
+		self.tx.send(msg).await.expect("test overseer no longer live");
+		Ok(())
+	}
+
+	async fn send_messages<T>(&mut self, msgs: T) -> SubsystemResult<()>
+		where T: IntoIterator<Item = AllMessages> + Send, T::IntoIter: Send
+	{
+		let mut iter = stream::iter(msgs.into_iter().map(Ok));
+		self.tx.send_all(&mut iter).await.expect("test overseer no longer live");
+
+		Ok(())
+	}
+}
+
+/// A handle for interacting with the subsystem context.
+pub struct TestSubsystemContextHandle<M> {
+	tx: SingleItemSink<FromOverseer<M>>,
+	rx: mpsc::UnboundedReceiver<AllMessages>,
+}
+
+impl<M> TestSubsystemContextHandle<M> {
+	/// Send a message or signal to the subsystem. This resolves at the point in time where the
+	/// subsystem has _read_ the message.
+	pub async fn send(&mut self, from_overseer: FromOverseer<M>) {
+		self.tx.send(from_overseer).await.expect("Test subsystem no longer live");
+	}
+
+	/// Receive the next message from the subsystem.
+	pub async fn recv(&mut self) -> AllMessages {
+		self.rx.next().await.expect("Test subsystem no longer live")
+	}
+}
+
+/// Make a test subsystem context.
+pub fn make_subsystem_context<M, S>(spawn: S)
+	-> (TestSubsystemContext<M, S>, TestSubsystemContextHandle<M>)
+{
+	let (overseer_tx, overseer_rx) = single_item_sink();
+	let (all_messages_tx, all_messages_rx) = mpsc::unbounded();
+
+	(
+		TestSubsystemContext {
+			tx: all_messages_tx,
+			rx: overseer_rx,
+			spawn,
+		},
+		TestSubsystemContextHandle {
+			tx: overseer_tx,
+			rx: all_messages_rx
+		},
+	)
+}