Improve test coverage of the Notifications protocol (#13033)

* Add handler and upgrade tests

* Add tests for `behaviour.rs`

* Apply review comments

* Update dependencies

* Apply suggestions from code review

Co-authored-by: Dmitry Markin <dmitry@markin.tech>

* Apply review comments

* Fix clippy

* Update mockall

* Apply review comment

---------

Co-authored-by: Dmitry Markin <dmitry@markin.tech>

substrate/client/consensus/common/Cargo.toml

@@ -18,7 +18,7 @@ futures = { version = "0.3.21", features = ["thread-pool"] }
 futures-timer = "3.0.1"
 libp2p = "0.50.0"
 log = "0.4.17"
-mockall = "0.11.2"
+mockall = "0.11.3"
 parking_lot = "0.12.1"
 serde = { version = "1.0", features = ["derive"] }
 thiserror = "1.0.30"

substrate/client/network/Cargo.toml

@@ -28,6 +28,7 @@ ip_network = "0.4.1"
 libp2p = { version = "0.50.0", features = ["dns", "identify", "kad", "macros", "mdns", "mplex", "noise", "ping", "tcp",  "tokio", "yamux", "websocket"] }
 log = "0.4.17"
 lru = "0.8.1"
+mockall = "0.11.3"
 parking_lot = "0.12.1"
 pin-project = "1.0.12"
 rand = "0.8.5"
@@ -52,9 +53,12 @@ sp-runtime = { version = "7.0.0", path = "../../primitives/runtime" }
 
 [dev-dependencies]
 assert_matches = "1.3"
+multistream-select = "0.12.1"
+rand = "0.8.5"
 tempfile = "3.1.0"
 tokio = { version = "1.22.0", features = ["macros"] }
 tokio-util = { version = "0.7.4", features = ["compat"] }
+tokio-test = "0.4.2"
 sc-network-light = { version = "0.10.0-dev", path = "./light" }
 sc-network-sync = { version = "0.10.0-dev", path = "./sync" }
 sp-test-primitives = { version = "2.0.0", path = "../../primitives/test-primitives" }

substrate/client/network/src/protocol/notifications/upgrade.rs

@@ -20,8 +20,8 @@ pub use self::{
 	collec::UpgradeCollec,
 	notifications::{
 		NotificationsHandshakeError, NotificationsIn, NotificationsInOpen,
-		NotificationsInSubstream, NotificationsOut, NotificationsOutError, NotificationsOutOpen,
+		NotificationsInSubstream, NotificationsInSubstreamHandshake, NotificationsOut,
-		NotificationsOutSubstream,
+		NotificationsOutError, NotificationsOutOpen, NotificationsOutSubstream,
 	},
 };
 

substrate/client/network/src/protocol/notifications/upgrade/collec.rs

@@ -73,7 +73,7 @@ where
 }
 
 /// Groups a `ProtocolName` with a `usize`.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq)]
 pub struct ProtoNameWithUsize<T>(T, usize);
 
 impl<T: ProtocolName> ProtocolName for ProtoNameWithUsize<T> {
@@ -99,3 +99,81 @@ impl<T: Future<Output = Result<O, E>>, O, E> Future for FutWithUsize<T> {
 		}
 	}
 }
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+	use libp2p::core::upgrade::{ProtocolName, UpgradeInfo};
+	use sc_network_common::protocol::ProtocolName as ProtoName;
+
+	// TODO: move to mocks
+	mockall::mock! {
+		pub ProtocolUpgrade<T> {}
+
+		impl<T: Clone + ProtocolName> UpgradeInfo for ProtocolUpgrade<T> {
+			type Info = T;
+			type InfoIter = vec::IntoIter<T>;
+			fn protocol_info(&self) -> vec::IntoIter<T>;
+		}
+	}
+
+	#[test]
+	fn protocol_info() {
+		let upgrades = (1..=3)
+			.map(|i| {
+				let mut upgrade = MockProtocolUpgrade::<ProtoNameWithUsize<ProtoName>>::new();
+				upgrade.expect_protocol_info().return_once(move || {
+					vec![ProtoNameWithUsize(ProtoName::from(format!("protocol{i}")), i)].into_iter()
+				});
+				upgrade
+			})
+			.collect::<Vec<_>>();
+
+		let upgrade: UpgradeCollec<_> = upgrades.into_iter().collect::<UpgradeCollec<_>>();
+		let protos = vec![
+			ProtoNameWithUsize(ProtoNameWithUsize(ProtoName::from("protocol1".to_string()), 1), 0),
+			ProtoNameWithUsize(ProtoNameWithUsize(ProtoName::from("protocol2".to_string()), 2), 1),
+			ProtoNameWithUsize(ProtoNameWithUsize(ProtoName::from("protocol3".to_string()), 3), 2),
+		];
+		let upgrades = upgrade.protocol_info().collect::<Vec<_>>();
+
+		assert_eq!(upgrades, protos,);
+	}
+
+	#[test]
+	fn nested_protocol_info() {
+		let mut upgrades = (1..=2)
+			.map(|i| {
+				let mut upgrade = MockProtocolUpgrade::<ProtoNameWithUsize<ProtoName>>::new();
+				upgrade.expect_protocol_info().return_once(move || {
+					vec![ProtoNameWithUsize(ProtoName::from(format!("protocol{i}")), i)].into_iter()
+				});
+				upgrade
+			})
+			.collect::<Vec<_>>();
+
+		upgrades.push({
+			let mut upgrade = MockProtocolUpgrade::<ProtoNameWithUsize<ProtoName>>::new();
+			upgrade.expect_protocol_info().return_once(move || {
+				vec![
+					ProtoNameWithUsize(ProtoName::from("protocol22".to_string()), 1),
+					ProtoNameWithUsize(ProtoName::from("protocol33".to_string()), 2),
+					ProtoNameWithUsize(ProtoName::from("protocol44".to_string()), 3),
+				]
+				.into_iter()
+			});
+			upgrade
+		});
+
+		let upgrade: UpgradeCollec<_> = upgrades.into_iter().collect::<UpgradeCollec<_>>();
+		let protos = vec![
+			ProtoNameWithUsize(ProtoNameWithUsize(ProtoName::from("protocol1".to_string()), 1), 0),
+			ProtoNameWithUsize(ProtoNameWithUsize(ProtoName::from("protocol2".to_string()), 2), 1),
+			ProtoNameWithUsize(ProtoNameWithUsize(ProtoName::from("protocol22".to_string()), 1), 2),
+			ProtoNameWithUsize(ProtoNameWithUsize(ProtoName::from("protocol33".to_string()), 2), 2),
+			ProtoNameWithUsize(ProtoNameWithUsize(ProtoName::from("protocol44".to_string()), 3), 2),
+		];
+		let upgrades = upgrade.protocol_info().collect::<Vec<_>>();
+		assert_eq!(upgrades, protos,);
+	}
+}

substrate/client/network/src/protocol/notifications/upgrade/notifications.rs

@@ -88,7 +88,8 @@ pub struct NotificationsInSubstream<TSubstream> {
 }
 
 /// State of the handshake sending back process.
-enum NotificationsInSubstreamHandshake {
+#[derive(Debug)]
+pub enum NotificationsInSubstreamHandshake {
 	/// Waiting for the user to give us the handshake message.
 	NotSent,
 	/// User gave us the handshake message. Trying to push it in the socket.
@@ -111,6 +112,13 @@ pub struct NotificationsOutSubstream<TSubstream> {
 	socket: Framed<TSubstream, UviBytes<io::Cursor<Vec<u8>>>>,
 }
 
+#[cfg(test)]
+impl<TSubstream> NotificationsOutSubstream<TSubstream> {
+	pub fn new(socket: Framed<TSubstream, UviBytes<io::Cursor<Vec<u8>>>>) -> Self {
+		Self { socket }
+	}
+}
+
 impl NotificationsIn {
 	/// Builds a new potential upgrade.
 	pub fn new(
@@ -193,6 +201,14 @@ impl<TSubstream> NotificationsInSubstream<TSubstream>
 where
 	TSubstream: AsyncRead + AsyncWrite + Unpin,
 {
+	#[cfg(test)]
+	pub fn new(
+		socket: Framed<TSubstream, UviBytes<io::Cursor<Vec<u8>>>>,
+		handshake: NotificationsInSubstreamHandshake,
+	) -> Self {
+		Self { socket, handshake }
+	}
+
 	/// Sends the handshake in order to inform the remote that we accept the substream.
 	pub fn send_handshake(&mut self, message: impl Into<Vec<u8>>) {
 		if !matches!(self.handshake, NotificationsInSubstreamHandshake::NotSent) {

substrate/client/network/sync/Cargo.toml

@@ -23,7 +23,7 @@ futures = "0.3.21"
 libp2p = "0.50.0"
 log = "0.4.17"
 lru = "0.8.1"
-mockall = "0.11.2"
+mockall = "0.11.3"
 prost = "0.11"
 smallvec = "1.8.0"
 thiserror = "1.0"