Fix block propagation between non-collator nodes (#363)

* Create builder for test nodes

* Fix syncing issue

* Adds missing file

cumulus/client/collator/src/lib.rs

@@ -79,7 +79,7 @@ where
 	fn new(
 		block_status: Arc<BS>,
 		spawner: Arc<dyn SpawnNamed + Send + Sync>,
-		announce_block: Arc<dyn Fn(Block::Hash, Vec<u8>) + Send + Sync>,
+		announce_block: Arc<dyn Fn(Block::Hash, Option<Vec<u8>>) + Send + Sync>,
 		backend: Arc<Backend>,
 		parachain_consensus: Box<dyn ParachainConsensus<Block>>,
 	) -> Self {
@@ -337,7 +337,7 @@ pub struct StartCollatorParams<Block: BlockT, Backend, BS, Spawner> {
 	pub para_id: ParaId,
 	pub backend: Arc<Backend>,
 	pub block_status: Arc<BS>,
-	pub announce_block: Arc<dyn Fn(Block::Hash, Vec<u8>) + Send + Sync>,
+	pub announce_block: Arc<dyn Fn(Block::Hash, Option<Vec<u8>>) + Send + Sync>,
 	pub overseer_handler: OverseerHandler,
 	pub spawner: Spawner,
 	pub key: CollatorPair,

cumulus/client/consensus/common/Cargo.toml

@@ -26,7 +26,7 @@ polkadot-runtime = { git = "https://github.com/paritytech/polkadot", branch = "m
 futures = { version = "0.3.8", features = ["compat"] }
 tokio = "0.1.22"
 codec = { package = "parity-scale-codec", version = "2.0.0", features = [ "derive" ] }
-tracing = "0.1.22"
+tracing = "0.1.25"
 async-trait = "0.1.42"
 dyn-clone = "1.0.4"
 

cumulus/client/consensus/common/src/lib.rs

@@ -143,7 +143,7 @@ pub async fn run_parachain_consensus<P, R, Block, B>(
 	para_id: ParaId,
 	parachain: Arc<P>,
 	relay_chain: R,
-	announce_block: Arc<dyn Fn(Block::Hash, Vec<u8>) + Send + Sync>,
+	announce_block: Arc<dyn Fn(Block::Hash, Option<Vec<u8>>) + Send + Sync>,
 ) -> ClientResult<()>
 where
 	Block: BlockT,
@@ -175,7 +175,7 @@ async fn follow_new_best<P, R, Block, B>(
 	para_id: ParaId,
 	parachain: Arc<P>,
 	relay_chain: R,
-	announce_block: Arc<dyn Fn(Block::Hash, Vec<u8>) + Send + Sync>,
+	announce_block: Arc<dyn Fn(Block::Hash, Option<Vec<u8>>) + Send + Sync>,
 ) -> ClientResult<()>
 where
 	Block: BlockT,
@@ -203,7 +203,6 @@ where
 					Some(h) => handle_new_best_parachain_head(
 						h,
 						&*parachain,
-						&*announce_block,
 						&mut unset_best_header,
 					),
 					None => {
@@ -241,12 +240,19 @@ fn handle_new_block_imported<Block, P>(
 	notification: BlockImportNotification<Block>,
 	unset_best_header_opt: &mut Option<Block::Header>,
 	parachain: &P,
-	announce_block: &dyn Fn(Block::Hash, Vec<u8>),
+	announce_block: &dyn Fn(Block::Hash, Option<Vec<u8>>),
 ) where
 	Block: BlockT,
 	P: UsageProvider<Block> + Send + Sync + BlockBackend<Block>,
 	for<'a> &'a P: BlockImport<Block>,
 {
+	// HACK
+	//
+	// Remove after https://github.com/paritytech/substrate/pull/8052 or similar is merged
+	if notification.origin != BlockOrigin::Own {
+		announce_block(notification.hash, None);
+	}
+
 	let unset_best_header = match (notification.is_new_best, &unset_best_header_opt) {
 		// If this is the new best block or we don't have any unset block, we can end it here.
 		(true, _) | (_, None) => return,
@@ -274,12 +280,12 @@ fn handle_new_block_imported<Block, P>(
 				.take()
 				.expect("We checked above that the value is set; qed");
 
-			import_block_as_new_best(unset_hash, unset_best_header, parachain, announce_block);
+			import_block_as_new_best(unset_hash, unset_best_header, parachain);
 		}
 		state => tracing::debug!(
 			target: "cumulus-consensus",
-			unset_best_header = ?unset_best_header,
-			imported_header = ?notification.header,
+			?unset_best_header,
+			?notification.header,
 			?state,
 			"Unexpected state for unset best header.",
 		),
@@ -290,7 +296,6 @@ fn handle_new_block_imported<Block, P>(
 fn handle_new_best_parachain_head<Block, P>(
 	head: Vec<u8>,
 	parachain: &P,
-	announce_block: &dyn Fn(Block::Hash, Vec<u8>),
 	unset_best_header: &mut Option<Block::Header>,
 ) where
 	Block: BlockT,
@@ -323,7 +328,7 @@ fn handle_new_best_parachain_head<Block, P>(
 			Ok(BlockStatus::InChainWithState) => {
 				unset_best_header.take();
 
-				import_block_as_new_best(hash, parachain_head, parachain, announce_block);
+				import_block_as_new_best(hash, parachain_head, parachain);
 			}
 			Ok(BlockStatus::InChainPruned) => {
 				tracing::error!(
@@ -358,7 +363,6 @@ fn import_block_as_new_best<Block, P>(
 	hash: Block::Hash,
 	header: Block::Header,
 	parachain: &P,
-	announce_block: &dyn Fn(Block::Hash, Vec<u8>),
 ) where
 	Block: BlockT,
 	P: UsageProvider<Block> + Send + Sync + BlockBackend<Block>,
@@ -376,8 +380,6 @@ fn import_block_as_new_best<Block, P>(
 			error = ?err,
 			"Failed to set new best block.",
 		);
-	} else {
-		(*announce_block)(hash, Vec::new());
 	}
 }
 

cumulus/client/network/Cargo.toml

@@ -30,11 +30,15 @@ parking_lot = "0.10.2"
 derive_more = "0.99.2"
 
 [dev-dependencies]
+tokio = { version = "0.2.21", features = ["macros"] }
+
 # Cumulus deps
 cumulus-test-service = { path = "../../test/service" }
+cumulus-primitives-core = { path = "../../primitives/core" }
 
 # Polkadot deps
 polkadot-test-client = { git = "https://github.com/paritytech/polkadot", branch = "master" }
+polkadot-test-service = { git = "https://github.com/paritytech/polkadot", branch = "master" }
 
 # substrate deps
 sc-client-api = { git = "https://github.com/paritytech/substrate", branch = "master" }
@@ -42,3 +46,6 @@ sp-core = { git = "https://github.com/paritytech/substrate", branch = "master" }
 sp-keyring = { git = "https://github.com/paritytech/substrate", branch = "master" }
 sp-keystore = { git = "https://github.com/paritytech/substrate", branch = "master" }
 sp-consensus = { git = "https://github.com/paritytech/substrate", branch = "master" }
+sc-cli = { git = "https://github.com/paritytech/substrate", branch = "master" }
+substrate-test-utils = { git = "https://github.com/paritytech/substrate", branch = "master" }
+sc-service = { git = "https://github.com/paritytech/substrate", branch = "master" }

cumulus/client/network/src/lib.rs

@@ -473,14 +473,14 @@ where
 /// the previous task running.
 pub struct WaitToAnnounce<Block: BlockT> {
 	spawner: Arc<dyn SpawnNamed + Send + Sync>,
-	announce_block: Arc<dyn Fn(Block::Hash, Vec<u8>) + Send + Sync>,
+	announce_block: Arc<dyn Fn(Block::Hash, Option<Vec<u8>>) + Send + Sync>,
 }
 
 impl<Block: BlockT> WaitToAnnounce<Block> {
 	/// Create the `WaitToAnnounce` object
 	pub fn new(
 		spawner: Arc<dyn SpawnNamed + Send + Sync>,
-		announce_block: Arc<dyn Fn(Block::Hash, Vec<u8>) + Send + Sync>,
+		announce_block: Arc<dyn Fn(Block::Hash, Option<Vec<u8>>) + Send + Sync>,
 	) -> WaitToAnnounce<Block> {
 		WaitToAnnounce {
 			spawner,
@@ -522,7 +522,7 @@ impl<Block: BlockT> WaitToAnnounce<Block> {
 async fn wait_to_announce<Block: BlockT>(
 	block_hash: <Block as BlockT>::Hash,
 	pov_hash: PHash,
-	announce_block: Arc<dyn Fn(Block::Hash, Vec<u8>) + Send + Sync>,
+	announce_block: Arc<dyn Fn(Block::Hash, Option<Vec<u8>>) + Send + Sync>,
 	signed_stmt_recv: oneshot::Receiver<SignedFullStatement>,
 ) {
 	let statement = match signed_stmt_recv.await {
@@ -541,7 +541,7 @@ async fn wait_to_announce<Block: BlockT>(
 	match statement.payload() {
 		Statement::Seconded(c) if &c.descriptor.pov_hash == &pov_hash => {
 			if let Ok(data) = BlockAnnounceData::try_from(statement) {
-				announce_block(block_hash, data.encode());
+				announce_block(block_hash, Some(data.encode()));
 			}
 		}
 		_ => tracing::debug!(

cumulus/client/network/tests/sync.rs

@@ -0,0 +1,86 @@
+// Copyright 2021 Parity Technologies (UK) Ltd.
+// This file is part of Substrate.
+
+// Substrate 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.
+
+// Substrate 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 Substrate.  If not, see <http://www.gnu.org/licenses/>.
+
+use cumulus_primitives_core::ParaId;
+use cumulus_test_service::{initial_head_data, Keyring::*};
+use futures::join;
+use sc_service::TaskExecutor;
+
+#[substrate_test_utils::test]
+async fn sync_blocks_from_tip_without_being_connected_to_a_collator(task_executor: TaskExecutor) {
+	let mut builder = sc_cli::LoggerBuilder::new("");
+	builder.with_colors(false);
+	let _ = builder.init();
+
+	let para_id = ParaId::from(100);
+
+	// start alice
+	let alice =
+		polkadot_test_service::run_validator_node(task_executor.clone(), Alice, || {}, vec![]);
+
+	// start bob
+	let bob = polkadot_test_service::run_validator_node(
+		task_executor.clone(),
+		Bob,
+		|| {},
+		vec![alice.addr.clone()],
+	);
+
+	// register parachain
+	alice
+		.register_parachain(
+			para_id,
+			cumulus_test_service::runtime::WASM_BINARY
+				.expect("You need to build the WASM binary to run this test!")
+				.to_vec(),
+			initial_head_data(para_id),
+		)
+		.await
+		.unwrap();
+
+	// run charlie as parachain collator
+	let charlie =
+		cumulus_test_service::TestNodeBuilder::new(para_id, task_executor.clone(), Charlie)
+			.enable_collator()
+			.connect_to_relay_chain_nodes(vec![&alice, &bob])
+			.build()
+			.await;
+
+	// run dave as parachain full node
+	let dave = cumulus_test_service::TestNodeBuilder::new(para_id, task_executor.clone(), Dave)
+		.connect_to_parachain_node(&charlie)
+		.connect_to_relay_chain_nodes(vec![&alice, &bob])
+		.build()
+		.await;
+
+	// run eve as parachain full node that is only connected to dave
+	let eve = cumulus_test_service::TestNodeBuilder::new(para_id, task_executor.clone(), Eve)
+		.connect_to_parachain_node(&dave)
+		.exclusively_connect_to_registered_parachain_nodes()
+		.connect_to_relay_chain_nodes(vec![&alice, &bob])
+		.build()
+		.await;
+
+	eve.wait_for_blocks(7).await;
+
+	join!(
+		alice.task_manager.clean_shutdown(),
+		bob.task_manager.clean_shutdown(),
+		charlie.task_manager.clean_shutdown(),
+		dave.task_manager.clean_shutdown(),
+		eve.task_manager.clean_shutdown(),
+	);
+}

cumulus/client/service/src/lib.rs

@@ -44,7 +44,7 @@ pub struct StartCollatorParams<'a, Block: BlockT, BS, Client, Backend, Spawner,
 	pub backend: Arc<Backend>,
 	pub block_status: Arc<BS>,
 	pub client: Arc<Client>,
-	pub announce_block: Arc<dyn Fn(Block::Hash, Vec<u8>) + Send + Sync>,
+	pub announce_block: Arc<dyn Fn(Block::Hash, Option<Vec<u8>>) + Send + Sync>,
 	pub spawner: Spawner,
 	pub para_id: ParaId,
 	pub collator_key: CollatorPair,
@@ -121,7 +121,7 @@ pub struct StartFullNodeParams<'a, Block: BlockT, Client, PClient> {
 	pub client: Arc<Client>,
 	pub polkadot_full_node: RFullNode<PClient>,
 	pub task_manager: &'a mut TaskManager,
-	pub announce_block: Arc<dyn Fn(Block::Hash, Vec<u8>) + Send + Sync>,
+	pub announce_block: Arc<dyn Fn(Block::Hash, Option<Vec<u8>>) + Send + Sync>,
 }
 
 /// Start a full node for a parachain.
@@ -165,7 +165,7 @@ where
 
 struct StartConsensus<'a, Block: BlockT, Client, Backend> {
 	para_id: ParaId,
-	announce_block: Arc<dyn Fn(Block::Hash, Vec<u8>) + Send + Sync>,
+	announce_block: Arc<dyn Fn(Block::Hash, Option<Vec<u8>>) + Send + Sync>,
 	client: Arc<Client>,
 	task_manager: &'a mut TaskManager,
 	_phantom: PhantomData<Backend>,

cumulus/rococo-parachains/src/service.rs

@@ -207,7 +207,7 @@ where
 
 	let announce_block = {
 		let network = network.clone();
-		Arc::new(move |hash, data| network.announce_block(hash, Some(data)))
+		Arc::new(move |hash, data| network.announce_block(hash, data))
 	};
 
 	if validator {

cumulus/test/service/Cargo.toml

@@ -56,12 +56,6 @@ jsonrpc-core = "15.1.0"
 futures = { version = "0.3.5" }
 tokio = { version = "0.2.21", features = ["macros"] }
 
-# Polkadot dependencies
-polkadot-test-runtime = { git = "https://github.com/paritytech/polkadot", branch = "master" }
-polkadot-test-service = { git = "https://github.com/paritytech/polkadot", branch = "master" }
-
 # Substrate dependencies
-pallet-sudo = { git = "https://github.com/paritytech/substrate", branch = "master" }
-substrate-test-runtime-client = { git = "https://github.com/paritytech/substrate", branch = "master" }
 substrate-test-utils = { git = "https://github.com/paritytech/substrate", branch = "master" }
 sc-cli = { git = "https://github.com/paritytech/substrate", branch = "master" }

cumulus/test/service/src/lib.rs

@@ -21,10 +21,6 @@
 mod chain_spec;
 mod genesis;
 
-pub use chain_spec::*;
-pub use cumulus_test_runtime as runtime;
-pub use genesis::*;
-
 use core::future::Future;
 use cumulus_client_network::BlockAnnounceValidator;
 use cumulus_client_service::{
@@ -53,6 +49,11 @@ use sp_trie::PrefixedMemoryDB;
 use std::sync::Arc;
 use substrate_test_client::BlockchainEventsExt;
 
+pub use chain_spec::*;
+pub use cumulus_test_runtime as runtime;
+pub use genesis::*;
+pub use sp_keyring::Sr25519Keyring as Keyring;
+
 // Native executor instance.
 native_executor_instance!(
 	pub RuntimeExecutor,
@@ -122,10 +123,9 @@ pub fn new_partial(
 #[sc_tracing::logging::prefix_logs_with(parachain_config.network.node_name.as_str())]
 async fn start_node_impl<RB>(
 	parachain_config: Configuration,
-	collator_key: CollatorPair,
+	collator_key: Option<CollatorPair>,
 	relay_chain_config: Configuration,
 	para_id: ParaId,
-	is_collator: bool,
 	rpc_ext_builder: RB,
 ) -> sc_service::error::Result<(
 	TaskManager,
@@ -157,7 +157,11 @@ where
 
 	let relay_chain_full_node = polkadot_test_service::new_full(
 		relay_chain_config,
-		polkadot_service::IsCollator::Yes(collator_key.public()),
+		if let Some(ref key) = collator_key {
+			polkadot_service::IsCollator::Yes(key.public())
+		} else {
+			polkadot_service::IsCollator::No
+		},
 	)
 	.map_err(|e| match e {
 		polkadot_service::Error::Sub(x) => x,
@@ -212,10 +216,10 @@ where
 
 	let announce_block = {
 		let network = network.clone();
-		Arc::new(move |hash, data| network.announce_block(hash, Some(data)))
+		Arc::new(move |hash, data| network.announce_block(hash, data))
 	};
 
-	if is_collator {
+	if let Some(collator_key) = collator_key {
 		let proposer_factory = sc_basic_authorship::ProposerFactory::with_proof_recording(
 			task_manager.spawn_handle(),
 			client.clone(),
@@ -270,7 +274,7 @@ where
 }
 
 /// A Cumulus test node instance used for testing.
-pub struct CumulusTestNode {
+pub struct TestNode {
 	/// TaskManager's instance.
 	pub task_manager: TaskManager,
 	/// Client's instance.
@@ -284,73 +288,156 @@ pub struct CumulusTestNode {
 	pub rpc_handlers: RpcHandlers,
 }
 
-/// Run a Cumulus test node using the Cumulus test runtime. The node will be using an in-memory
-/// socket, therefore you need to provide boot nodes if you want it to be connected to other nodes.
-/// The `storage_update_func` can be used to make adjustements to the runtime before the node
-/// starts.
-pub async fn run_test_node(
+
+/// A builder to create a [`TestNode`].
+pub struct TestNodeBuilder {
+	para_id: ParaId,
 	task_executor: TaskExecutor,
 	key: Sr25519Keyring,
-	parachain_storage_update_func: impl Fn(),
-	relay_chain_storage_update_func: impl Fn(),
-	parachain_boot_nodes: Vec<MultiaddrWithPeerId>,
-	relay_chain_boot_nodes: Vec<MultiaddrWithPeerId>,
-	para_id: ParaId,
-	is_collator: bool,
-) -> CumulusTestNode {
-	let collator_key = CollatorPair::generate().0;
-	let parachain_config = node_config(
-		parachain_storage_update_func,
-		task_executor.clone(),
-		key,
-		parachain_boot_nodes,
-		para_id,
-		is_collator,
-	)
-	.expect("could not generate Configuration");
-	let mut relay_chain_config = polkadot_test_service::node_config(
-		relay_chain_storage_update_func,
-		task_executor.clone(),
-		key,
-		relay_chain_boot_nodes,
-		false,
-	);
+	collator_key: Option<CollatorPair>,
+	parachain_nodes: Vec<MultiaddrWithPeerId>,
+	parachain_nodes_exclusive: bool,
+	relay_chain_nodes: Vec<MultiaddrWithPeerId>,
+}
 
-	relay_chain_config.network.node_name =
-		format!("{} (relay chain)", relay_chain_config.network.node_name);
+impl TestNodeBuilder {
+	/// Create a new instance of `Self`.
+	///
+	/// `para_id` - The parachain id this node is running for.
+	/// `task_executor` - The task executor to use.
+	/// `key` - The key that will be used to generate the name and that will be passed as `dev_seed`.
+	pub fn new(para_id: ParaId, task_executor: TaskExecutor, key: Sr25519Keyring) -> Self {
+		TestNodeBuilder {
+			key,
+			para_id,
+			task_executor,
+			collator_key: None,
+			parachain_nodes: Vec::new(),
+			parachain_nodes_exclusive: false,
+			relay_chain_nodes: Vec::new(),
+		}
+	}
 
-	let multiaddr = parachain_config.network.listen_addresses[0].clone();
-	let (task_manager, client, network, rpc_handlers) = start_node_impl(
-		parachain_config,
-		collator_key,
-		relay_chain_config,
-		para_id,
-		is_collator,
-		|_| Default::default(),
-	)
-	.await
-	.expect("could not create Cumulus test service");
+	/// Enable collator for this node.
+	pub fn enable_collator(mut self) -> Self {
+		let collator_key = CollatorPair::generate().0;
+		self.collator_key = Some(collator_key);
+		self
+	}
 
-	let peer_id = network.local_peer_id().clone();
-	let addr = MultiaddrWithPeerId { multiaddr, peer_id };
+	/// Instruct the node to exclusively connect to registered parachain nodes.
+	///
+	/// Parachain nodes can be registered using [`Self::connect_to_parachain_node`] and
+	/// [`Self::connect_to_parachain_nodes`].
+	pub fn exclusively_connect_to_registered_parachain_nodes(mut self) -> Self {
+		self.parachain_nodes_exclusive = true;
+		self
+	}
 
-	CumulusTestNode {
-		task_manager,
-		client,
-		network,
-		addr,
-		rpc_handlers,
+	/// Make the node connect to the given parachain node.
+	///
+	/// By default the node will not be connected to any node or will be able to discover any other
+	/// node.
+	pub fn connect_to_parachain_node(mut self, node: &TestNode) -> Self {
+		self.parachain_nodes.push(node.addr.clone());
+		self
+	}
+
+	/// Make the node connect to the given parachain nodes.
+	///
+	/// By default the node will not be connected to any node or will be able to discover any other
+	/// node.
+	pub fn connect_to_parachain_nodes<'a>(
+		mut self,
+		nodes: impl Iterator<Item = &'a TestNode>,
+	) -> Self {
+		self.parachain_nodes.extend(nodes.map(|n| n.addr.clone()));
+		self
+	}
+
+	/// Make the node connect to the given relay chain node.
+	///
+	/// By default the node will not be connected to any node or will be able to discover any other
+	/// node.
+	pub fn connect_to_relay_chain_node(
+		mut self,
+		node: &polkadot_test_service::PolkadotTestNode,
+	) -> Self {
+		self.relay_chain_nodes.push(node.addr.clone());
+		self
+	}
+
+	/// Make the node connect to the given relay chain nodes.
+	///
+	/// By default the node will not be connected to any node or will be able to discover any other
+	/// node.
+	pub fn connect_to_relay_chain_nodes<'a>(
+		mut self,
+		nodes: impl IntoIterator<Item = &'a polkadot_test_service::PolkadotTestNode>,
+	) -> Self {
+		self.relay_chain_nodes.extend(nodes.into_iter().map(|n| n.addr.clone()));
+		self
+	}
+
+	/// Build the [`TestNode`].
+	pub async fn build(self) -> TestNode {
+		let parachain_config = node_config(
+			|| (),
+			self.task_executor.clone(),
+			self.key.clone(),
+			self.parachain_nodes,
+			self.parachain_nodes_exclusive,
+			self.para_id,
+			self.collator_key.is_some(),
+		)
+		.expect("could not generate Configuration");
+		let mut relay_chain_config = polkadot_test_service::node_config(
+			|| (),
+			self.task_executor,
+			self.key,
+			self.relay_chain_nodes,
+			false,
+		);
+
+		relay_chain_config.network.node_name =
+			format!("{} (relay chain)", relay_chain_config.network.node_name);
+
+		let multiaddr = parachain_config.network.listen_addresses[0].clone();
+		let (task_manager, client, network, rpc_handlers) = start_node_impl(
+			parachain_config,
+			self.collator_key,
+			relay_chain_config,
+			self.para_id,
+			|_| Default::default(),
+		)
+		.await
+		.expect("could not create Cumulus test service");
+
+		let peer_id = network.local_peer_id().clone();
+		let addr = MultiaddrWithPeerId { multiaddr, peer_id };
+
+		TestNode {
+			task_manager,
+			client,
+			network,
+			addr,
+			rpc_handlers,
+		}
 	}
 }
 
-/// Create a Cumulus `Configuration`. By default an in-memory socket will be used, therefore you
-/// need to provide boot nodes if you want the future node to be connected to other nodes. The
-/// `storage_update_func` can be used to make adjustments to the runtime before the node starts.
+/// Create a Cumulus `Configuration`.
+///
+/// By default an in-memory socket will be used, therefore you need to provide nodes if you want the
+/// node to be connected to other nodes. If `nodes_exclusive` is `true`, the node will only connect
+/// to the given `nodes` and not to any other node. The `storage_update_func` can be used to make
+/// adjustments to the runtime genesis.
 pub fn node_config(
 	storage_update_func: impl Fn(),
 	task_executor: TaskExecutor,
 	key: Sr25519Keyring,
-	boot_nodes: Vec<MultiaddrWithPeerId>,
+	nodes: Vec<MultiaddrWithPeerId>,
+	nodes_exlusive: bool,
 	para_id: ParaId,
 	is_collator: bool,
 ) -> Result<Configuration, ServiceError> {
@@ -379,7 +466,13 @@ pub fn node_config(
 		None,
 	);
 
-	network_config.boot_nodes = boot_nodes;
+	if nodes_exlusive {
+		network_config.default_peers_set.reserved_nodes = nodes;
+		network_config.default_peers_set.non_reserved_mode =
+			sc_network::config::NonReservedPeerMode::Deny;
+	} else {
+		network_config.boot_nodes = nodes;
+	}
 
 	network_config.allow_non_globals_in_dht = true;
 
@@ -445,7 +538,7 @@ pub fn node_config(
 	})
 }
 
-impl CumulusTestNode {
+impl TestNode {
 	/// Wait for `count` blocks to be imported in the node and then exit. This function will not
 	/// return if no blocks are ever created, thus you should restrict the maximum amount of time of
 	/// the test execution.

cumulus/test/service/tests/integration.rs

@@ -15,10 +15,9 @@
 // along with Substrate.  If not, see <http://www.gnu.org/licenses/>.
 
 use cumulus_primitives_core::ParaId;
-use cumulus_test_service::initial_head_data;
+use cumulus_test_service::{initial_head_data, Keyring::*};
 use futures::join;
 use sc_service::TaskExecutor;
-use substrate_test_runtime_client::AccountKeyring::*;
 
 #[substrate_test_utils::test]
 async fn test_collating_and_non_collator_mode_catching_up(task_executor: TaskExecutor) {
@@ -53,31 +52,21 @@ async fn test_collating_and_non_collator_mode_catching_up(task_executor: TaskExe
 		.unwrap();
 
 	// run cumulus charlie (a parachain collator)
-	let charlie = cumulus_test_service::run_test_node(
-		task_executor.clone(),
-		Charlie,
-		|| {},
-		|| {},
-		vec![],
-		vec![alice.addr.clone(), bob.addr.clone()],
-		para_id,
-		true,
-	)
-	.await;
+	let charlie =
+		cumulus_test_service::TestNodeBuilder::new(para_id, task_executor.clone(), Charlie)
+			.enable_collator()
+			.connect_to_relay_chain_nodes(vec![&alice, &bob])
+			.build()
+			.await;
 	charlie.wait_for_blocks(5).await;
 
 	// run cumulus dave (a parachain full node) and wait for it to sync some blocks
-	let dave = cumulus_test_service::run_test_node(
-		task_executor.clone(),
-		Dave,
-		|| {},
-		|| {},
-		vec![charlie.addr.clone()],
-		vec![alice.addr.clone(), bob.addr.clone()],
-		para_id,
-		false,
-	)
-	.await;
+	let dave = cumulus_test_service::TestNodeBuilder::new(para_id, task_executor.clone(), Dave)
+		.connect_to_parachain_node(&charlie)
+		.connect_to_relay_chain_nodes(vec![&alice, &bob])
+		.build()
+		.await;
+
 	dave.wait_for_blocks(7).await;
 
 	join!(