Less sleeps (#9848)

* Less sleeps

* No need to use tokio-test crate

* Less sleep

* Avoid leaving zombie substrates around

(when panicing in tests)

* Remove unused imports

* Incorporating feedback

* rename method

* Use rpc_api

* Update bin/node/cli/tests/temp_base_path_works.rs

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

* Update bin/node/cli/tests/temp_base_path_works.rs

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

substrate/bin/node/cli/tests/common.rs

@@ -23,58 +23,115 @@ use nix::{
 	sys::signal::{kill, Signal::SIGINT},
 	unistd::Pid,
 };
+use node_primitives::Block;
+use remote_externalities::rpc_api;
 use std::{
 	convert::TryInto,
+	ops::{Deref, DerefMut},
 	path::Path,
 	process::{Child, Command, ExitStatus},
-	thread,
 	time::Duration,
 };
+use tokio::time::timeout;
+
+static LOCALHOST_WS: &str = "ws://127.0.0.1:9944/";
 
 /// Wait for the given `child` the given number of `secs`.
 ///
 /// Returns the `Some(exit status)` or `None` if the process did not finish in the given time.
-pub fn wait_for(child: &mut Child, secs: usize) -> Option<ExitStatus> {
-	for i in 0..secs {
-		match child.try_wait().unwrap() {
-			Some(status) => {
-				if i > 5 {
-					eprintln!("Child process took {} seconds to exit gracefully", i);
-				}
-				return Some(status)
-			},
-			None => thread::sleep(Duration::from_secs(1)),
+pub fn wait_for(child: &mut Child, secs: u64) -> Result<ExitStatus, ()> {
+	let result = wait_timeout::ChildExt::wait_timeout(child, Duration::from_secs(5.min(secs)))
+		.map_err(|_| ())?;
+	if let Some(exit_status) = result {
+		Ok(exit_status)
+	} else {
+		if secs > 5 {
+			eprintln!("Child process taking over 5 seconds to exit gracefully");
+			let result = wait_timeout::ChildExt::wait_timeout(child, Duration::from_secs(secs - 5))
+				.map_err(|_| ())?;
+			if let Some(exit_status) = result {
+				return Ok(exit_status)
+			}
 		}
+		eprintln!("Took too long to exit (> {} seconds). Killing...", secs);
+		let _ = child.kill();
+		child.wait().unwrap();
+		Err(())
 	}
-	eprintln!("Took too long to exit (> {} seconds). Killing...", secs);
-	let _ = child.kill();
-	child.wait().unwrap();
-
-	None
 }
 
-/// Run the node for a while (30 seconds)
-pub fn run_node_for_a_while(base_path: &Path, args: &[&str]) {
+/// Wait for at least n blocks to be finalized within a specified time.
+pub async fn wait_n_finalized_blocks(
+	n: usize,
+	timeout_secs: u64,
+) -> Result<(), tokio::time::error::Elapsed> {
+	timeout(Duration::from_secs(timeout_secs), wait_n_finalized_blocks_from(n, LOCALHOST_WS)).await
+}
+
+/// Wait for at least n blocks to be finalized from a specified node
+pub async fn wait_n_finalized_blocks_from(n: usize, url: &str) {
+	let mut built_blocks = std::collections::HashSet::new();
+	let mut interval = tokio::time::interval(Duration::from_secs(2));
+
+	loop {
+		if let Ok(block) = rpc_api::get_finalized_head::<Block, _>(url.to_string()).await {
+			built_blocks.insert(block);
+			if built_blocks.len() > n {
+				break
+			}
+		};
+		interval.tick().await;
+	}
+}
+
+/// Run the node for a while (3 blocks)
+pub async fn run_node_for_a_while(base_path: &Path, args: &[&str]) {
 	let mut cmd = Command::new(cargo_bin("substrate"));
 
-	let mut cmd = cmd.args(args).arg("-d").arg(base_path).spawn().unwrap();
+	let mut child = KillChildOnDrop(cmd.args(args).arg("-d").arg(base_path).spawn().unwrap());
 
 	// Let it produce some blocks.
-	thread::sleep(Duration::from_secs(30));
-	assert!(cmd.try_wait().unwrap().is_none(), "the process should still be running");
+	let _ = wait_n_finalized_blocks(3, 30).await;
+
+	assert!(child.try_wait().unwrap().is_none(), "the process should still be running");
 
 	// Stop the process
-	kill(Pid::from_raw(cmd.id().try_into().unwrap()), SIGINT).unwrap();
-	assert!(wait_for(&mut cmd, 40).map(|x| x.success()).unwrap_or_default());
+	kill(Pid::from_raw(child.id().try_into().unwrap()), SIGINT).unwrap();
+	assert!(wait_for(&mut child, 40).map(|x| x.success()).unwrap());
 }
 
 /// Run the node asserting that it fails with an error
 pub fn run_node_assert_fail(base_path: &Path, args: &[&str]) {
 	let mut cmd = Command::new(cargo_bin("substrate"));
 
-	let mut cmd = cmd.args(args).arg("-d").arg(base_path).spawn().unwrap();
+	let mut child = KillChildOnDrop(cmd.args(args).arg("-d").arg(base_path).spawn().unwrap());
 
-	// Let it produce some blocks.
-	thread::sleep(Duration::from_secs(10));
-	assert!(cmd.try_wait().unwrap().is_some(), "the process should not be running anymore");
+	// Let it produce some blocks, but it should die within 10 seconds.
+	assert_ne!(
+		wait_timeout::ChildExt::wait_timeout(&mut *child, Duration::from_secs(10)).unwrap(),
+		None,
+		"the process should not be running anymore"
+	);
+}
+
+pub struct KillChildOnDrop(pub Child);
+
+impl Drop for KillChildOnDrop {
+	fn drop(&mut self) {
+		let _ = self.0.kill();
+	}
+}
+
+impl Deref for KillChildOnDrop {
+	type Target = Child;
+
+	fn deref(&self) -> &Self::Target {
+		&self.0
+	}
+}
+
+impl DerefMut for KillChildOnDrop {
+	fn deref_mut(&mut self) -> &mut Self::Target {
+		&mut self.0
+	}
 }