metrics: tests: Fix flaky runtime_can_publish_metrics

node/metrics/src/tests.rs

@@ -61,8 +61,8 @@ async fn runtime_can_publish_metrics() {
 	// Start validator Bob.
 	let _bob = run_validator_node(bob_config, None);
 
-	// Wait for Alice to author two blocks.
-	alice.wait_for_blocks(2).await;
+	// Wait for Alice to see two finalized blocks.
+	alice.wait_for_finalized_blocks(2).await;
 
 	let metrics_uri = format!("http://localhost:{}/metrics", DEFAULT_PROMETHEUS_PORT);
 	let metrics = scrape_prometheus_metrics(&metrics_uri).await;

node/test/service/src/lib.rs

@@ -21,7 +21,7 @@
 pub mod chain_spec;
 
 pub use chain_spec::*;
-use futures::future::Future;
+use futures::{future::Future, stream::StreamExt};
 use polkadot_node_primitives::{CollationGenerationConfig, CollatorFn};
 use polkadot_node_subsystem::messages::{CollationGenerationMessage, CollatorProtocolMessage};
 use polkadot_overseer::Handle;
@@ -35,8 +35,9 @@ use polkadot_test_runtime::{
 	ParasCall, ParasSudoWrapperCall, Runtime, SignedExtra, SignedPayload, SudoCall,
 	UncheckedExtrinsic, VERSION,
 };
+
 use sc_chain_spec::ChainSpec;
-use sc_client_api::execution_extensions::ExecutionStrategies;
+use sc_client_api::{execution_extensions::ExecutionStrategies, BlockchainEvents};
 use sc_network::{
 	config::{NetworkConfiguration, TransportConfig},
 	multiaddr, NetworkStateInfo,
@@ -54,14 +55,14 @@ use sp_keyring::Sr25519Keyring;
 use sp_runtime::{codec::Encode, generic, traits::IdentifyAccount, MultiSigner};
 use sp_state_machine::BasicExternalities;
 use std::{
+	collections::HashSet,
 	net::{Ipv4Addr, SocketAddr},
 	path::PathBuf,
 	sync::Arc,
 };
 use substrate_test_client::{
 	BlockchainEventsExt, RpcHandlersExt, RpcTransactionError, RpcTransactionOutput,
 };
-
 /// Declare an instance of the native executor named `PolkadotTestExecutorDispatch`. Include the wasm binary as the
 /// equivalent wasm code.
 pub struct PolkadotTestExecutorDispatch;
@@ -335,6 +336,23 @@ impl PolkadotTestNode {
 		self.client.wait_for_blocks(count)
 	}
 
+	/// Wait for `count` blocks to be finalized and then exit. Similarly with `wait_for_blocks` this function will
+	/// not return if no block are ever finalized.
+	pub async fn wait_for_finalized_blocks(&self, count: usize) {
+		let mut import_notification_stream = self.client.finality_notification_stream();
+		let mut blocks = HashSet::new();
+
+		Box::pin(async move {
+			while let Some(notification) = import_notification_stream.next().await {
+				blocks.insert(notification.hash);
+				if blocks.len() == count {
+					break
+				}
+			}
+		})
+		.await;
+	}
+
 	/// Register the collator functionality in the overseer of this node.
 	pub async fn register_collator(
 		&mut self,