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/network/src/protocol/notifications/upgrade.rs

@@ -20,8 +20,8 @@ pub use self::{
 	collec::UpgradeCollec,
 	notifications::{
 		NotificationsHandshakeError, NotificationsIn, NotificationsInOpen,
-		NotificationsInSubstream, NotificationsOut, NotificationsOutError, NotificationsOutOpen,
-		NotificationsOutSubstream,
+		NotificationsInSubstream, NotificationsInSubstreamHandshake, NotificationsOut,
+		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) {