referrerpolicy=no-referrer-when-downgrade

cumulus_test_service/
lib.rs

1// Copyright (C) Parity Technologies (UK) Ltd.
2// This file is part of Cumulus.
3
4// Cumulus is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, either version 3 of the License, or
7// (at your option) any later version.
8
9// Cumulus is distributed in the hope that it will be useful,
10// but WITHOUT ANY WARRANTY; without even the implied warranty of
11// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12// GNU General Public License for more details.
13
14// You should have received a copy of the GNU General Public License
15// along with Cumulus.  If not, see <http://www.gnu.org/licenses/>.
16
17//! Crate used for testing with Cumulus.
18
19#![warn(missing_docs)]
20
21/// Utilities used for benchmarking
22pub mod bench_utils;
23
24pub mod chain_spec;
25
26use cumulus_client_collator::service::CollatorService;
27use cumulus_client_consensus_aura::{
28	collators::{
29		lookahead::{self as aura, Params as AuraParams},
30		slot_based::{
31			self as slot_based, Params as SlotBasedParams, SlotBasedBlockImport,
32			SlotBasedBlockImportHandle,
33		},
34	},
35	ImportQueueParams,
36};
37use prometheus::Registry;
38use runtime::AccountId;
39use sc_executor::{HeapAllocStrategy, WasmExecutor, DEFAULT_HEAP_ALLOC_STRATEGY};
40use sp_consensus_aura::sr25519::AuthorityPair;
41use std::{
42	collections::HashSet,
43	future::Future,
44	net::{Ipv4Addr, SocketAddr, SocketAddrV4},
45	time::Duration,
46};
47use url::Url;
48
49use crate::runtime::Weight;
50use cumulus_client_cli::{CollatorOptions, RelayChainMode};
51use cumulus_client_consensus_common::ParachainBlockImport as TParachainBlockImport;
52use cumulus_client_pov_recovery::{RecoveryDelayRange, RecoveryHandle};
53use cumulus_client_service::{
54	build_network, prepare_node_config, start_relay_chain_tasks, BuildNetworkParams,
55	CollatorSybilResistance, DARecoveryProfile, ParachainTracingExecuteBlock,
56	StartRelayChainTasksParams,
57};
58use cumulus_primitives_core::{relay_chain::ValidationCode, GetParachainInfo, ParaId};
59use cumulus_relay_chain_inprocess_interface::RelayChainInProcessInterface;
60use cumulus_relay_chain_interface::{RelayChainError, RelayChainInterface, RelayChainResult};
61use cumulus_relay_chain_minimal_node::build_minimal_relay_chain_node_with_rpc;
62
63use cumulus_test_runtime::{Hash, NodeBlock as Block, RuntimeApi};
64
65use frame_system_rpc_runtime_api::AccountNonceApi;
66use polkadot_node_subsystem::{errors::RecoveryError, messages::AvailabilityRecoveryMessage};
67use polkadot_overseer::Handle as OverseerHandle;
68use polkadot_primitives::{CandidateHash, CollatorPair};
69use polkadot_service::ProvideRuntimeApi;
70use sc_consensus::ImportQueue;
71use sc_network::{
72	config::{FullNetworkConfiguration, TransportConfig},
73	multiaddr,
74	service::traits::NetworkService,
75	NetworkBackend, NetworkBlock, NetworkStateInfo,
76};
77use sc_service::{
78	config::{
79		BlocksPruning, DatabaseSource, ExecutorConfiguration, KeystoreConfig, MultiaddrWithPeerId,
80		NetworkConfiguration, OffchainWorkerConfig, PruningMode, RpcBatchRequestConfig,
81		RpcConfiguration, RpcEndpoint, WasmExecutionMethod,
82	},
83	BasePath, ChainSpec as ChainSpecService, Configuration, Error as ServiceError,
84	PartialComponents, Role, RpcHandlers, TFullBackend, TFullClient, TaskManager,
85};
86use sp_arithmetic::traits::SaturatedConversion;
87use sp_blockchain::HeaderBackend;
88use sp_core::Pair;
89use sp_keyring::Sr25519Keyring;
90use sp_runtime::{codec::Encode, generic, MultiAddress};
91use sp_state_machine::BasicExternalities;
92use std::sync::Arc;
93use substrate_test_client::{
94	BlockchainEventsExt, RpcHandlersExt, RpcTransactionError, RpcTransactionOutput,
95};
96
97pub use chain_spec::*;
98pub use cumulus_test_runtime as runtime;
99pub use sp_keyring::Sr25519Keyring as Keyring;
100
101const LOG_TARGET: &str = "cumulus-test-service";
102
103/// The signature of the announce block fn.
104pub type AnnounceBlockFn = Arc<dyn Fn(Hash, Option<Vec<u8>>) + Send + Sync>;
105
106type HostFunctions =
107	(sp_io::SubstrateHostFunctions, cumulus_client_service::storage_proof_size::HostFunctions);
108/// The client type being used by the test service.
109pub type Client = TFullClient<runtime::NodeBlock, runtime::RuntimeApi, WasmExecutor<HostFunctions>>;
110
111/// The backend type being used by the test service.
112pub type Backend = TFullBackend<Block>;
113
114/// The block-import type being used by the test service.
115pub type ParachainBlockImport =
116	TParachainBlockImport<Block, SlotBasedBlockImport<Block, Arc<Client>, Client>, Backend>;
117
118/// Transaction pool type used by the test service
119pub type TransactionPool = Arc<sc_transaction_pool::TransactionPoolHandle<Block, Client>>;
120
121/// Recovery handle that fails regularly to simulate unavailable povs.
122pub struct FailingRecoveryHandle {
123	overseer_handle: OverseerHandle,
124	counter: u32,
125	failed_hashes: HashSet<CandidateHash>,
126}
127
128impl FailingRecoveryHandle {
129	/// Create a new FailingRecoveryHandle
130	pub fn new(overseer_handle: OverseerHandle) -> Self {
131		Self { overseer_handle, counter: 0, failed_hashes: Default::default() }
132	}
133}
134
135#[async_trait::async_trait]
136impl RecoveryHandle for FailingRecoveryHandle {
137	async fn send_recovery_msg(
138		&mut self,
139		message: AvailabilityRecoveryMessage,
140		origin: &'static str,
141	) {
142		let AvailabilityRecoveryMessage::RecoverAvailableData(ref receipt, _, _, _, _) = message;
143		let candidate_hash = receipt.hash();
144
145		// For every 3rd block we immediately signal unavailability to trigger
146		// a retry. The same candidate is never failed multiple times to ensure progress.
147		if self.counter % 3 == 0 && self.failed_hashes.insert(candidate_hash) {
148			tracing::info!(target: LOG_TARGET, ?candidate_hash, "Failing pov recovery.");
149
150			let AvailabilityRecoveryMessage::RecoverAvailableData(_, _, _, _, back_sender) =
151				message;
152			back_sender
153				.send(Err(RecoveryError::Unavailable))
154				.expect("Return channel should work here.");
155		} else {
156			self.overseer_handle.send_msg(message, origin).await;
157		}
158		self.counter += 1;
159	}
160}
161
162/// Assembly of PartialComponents (enough to run chain ops subcommands)
163pub type Service = PartialComponents<
164	Client,
165	Backend,
166	(),
167	sc_consensus::import_queue::BasicQueue<Block>,
168	sc_transaction_pool::TransactionPoolHandle<Block, Client>,
169	(ParachainBlockImport, SlotBasedBlockImportHandle<Block>),
170>;
171
172/// Starts a `ServiceBuilder` for a full service.
173///
174/// Use this macro if you don't actually need the full service, but just the builder in order to
175/// be able to perform chain operations.
176pub fn new_partial(
177	config: &mut Configuration,
178	enable_import_proof_record: bool,
179) -> Result<Service, sc_service::Error> {
180	let heap_pages = config
181		.executor
182		.default_heap_pages
183		.map_or(DEFAULT_HEAP_ALLOC_STRATEGY, |h| HeapAllocStrategy::Static { extra_pages: h as _ });
184
185	let executor = WasmExecutor::builder()
186		.with_execution_method(config.executor.wasm_method)
187		.with_onchain_heap_alloc_strategy(heap_pages)
188		.with_offchain_heap_alloc_strategy(heap_pages)
189		.with_max_runtime_instances(config.executor.max_runtime_instances)
190		.with_runtime_cache_size(config.executor.runtime_cache_size)
191		.build();
192
193	let (client, backend, keystore_container, task_manager) =
194		sc_service::new_full_parts_record_import::<Block, RuntimeApi, _>(
195			config,
196			None,
197			executor,
198			enable_import_proof_record,
199		)?;
200	let client = Arc::new(client);
201
202	let (block_import, slot_based_handle) =
203		SlotBasedBlockImport::new(client.clone(), client.clone());
204	let block_import = ParachainBlockImport::new(block_import, backend.clone());
205
206	let transaction_pool = Arc::from(
207		sc_transaction_pool::Builder::new(
208			task_manager.spawn_essential_handle(),
209			client.clone(),
210			config.role.is_authority().into(),
211		)
212		.with_options(config.transaction_pool.clone())
213		.with_prometheus(config.prometheus_registry())
214		.build(),
215	);
216
217	let slot_duration = sc_consensus_aura::slot_duration(&*client)?;
218	let import_queue = cumulus_client_consensus_aura::import_queue::<AuthorityPair, _, _, _, _, _>(
219		ImportQueueParams {
220			block_import: block_import.clone(),
221			client: client.clone(),
222			create_inherent_data_providers: move |_, ()| async move {
223				let timestamp = sp_timestamp::InherentDataProvider::from_system_time();
224
225				let slot =
226					sp_consensus_aura::inherents::InherentDataProvider::from_timestamp_and_slot_duration(
227						*timestamp,
228						slot_duration,
229					);
230
231				Ok((slot, timestamp))
232			},
233			spawner: &task_manager.spawn_essential_handle(),
234			registry: None,
235			telemetry: None,
236		},
237	)?;
238
239	let params = PartialComponents {
240		backend,
241		client,
242		import_queue,
243		keystore_container,
244		task_manager,
245		transaction_pool,
246		select_chain: (),
247		other: (block_import, slot_based_handle),
248	};
249
250	Ok(params)
251}
252
253async fn build_relay_chain_interface(
254	relay_chain_config: Configuration,
255	parachain_prometheus_registry: Option<&Registry>,
256	collator_key: Option<CollatorPair>,
257	collator_options: CollatorOptions,
258	task_manager: &mut TaskManager,
259) -> RelayChainResult<Arc<dyn RelayChainInterface + 'static>> {
260	let relay_chain_node = match collator_options.relay_chain_mode {
261		cumulus_client_cli::RelayChainMode::Embedded => polkadot_test_service::new_full(
262			relay_chain_config,
263			if let Some(ref key) = collator_key {
264				polkadot_service::IsParachainNode::Collator(key.clone())
265			} else {
266				polkadot_service::IsParachainNode::Collator(CollatorPair::generate().0)
267			},
268			None,
269			polkadot_service::CollatorOverseerGen,
270			Some("Relaychain"),
271		)
272		.map_err(|e| RelayChainError::Application(Box::new(e) as Box<_>))?,
273		cumulus_client_cli::RelayChainMode::ExternalRpc(rpc_target_urls) =>
274			return build_minimal_relay_chain_node_with_rpc(
275				relay_chain_config,
276				parachain_prometheus_registry,
277				task_manager,
278				rpc_target_urls,
279			)
280			.await
281			.map(|r| r.0),
282	};
283
284	task_manager.add_child(relay_chain_node.task_manager);
285	tracing::info!("Using inprocess node.");
286	Ok(Arc::new(RelayChainInProcessInterface::new(
287		relay_chain_node.client.clone(),
288		relay_chain_node.backend.clone(),
289		relay_chain_node.sync_service.clone(),
290		relay_chain_node.overseer_handle.ok_or(RelayChainError::GenericError(
291			"Overseer should be running in full node.".to_string(),
292		))?,
293	)))
294}
295
296/// Start a node with the given parachain `Configuration` and relay chain `Configuration`.
297///
298/// This is the actual implementation that is abstract over the executor and the runtime api.
299#[sc_tracing::logging::prefix_logs_with("Parachain")]
300pub async fn start_node_impl<RB, Net: NetworkBackend<Block, Hash>>(
301	parachain_config: Configuration,
302	collator_key: Option<CollatorPair>,
303	relay_chain_config: Configuration,
304	wrap_announce_block: Option<Box<dyn FnOnce(AnnounceBlockFn) -> AnnounceBlockFn>>,
305	fail_pov_recovery: bool,
306	rpc_ext_builder: RB,
307	collator_options: CollatorOptions,
308	proof_recording_during_import: bool,
309	use_slot_based_collator: bool,
310) -> sc_service::error::Result<(
311	TaskManager,
312	Arc<Client>,
313	Arc<dyn NetworkService>,
314	RpcHandlers,
315	TransactionPool,
316	Arc<Backend>,
317)>
318where
319	RB: Fn(Arc<Client>) -> Result<jsonrpsee::RpcModule<()>, sc_service::Error> + Send + 'static,
320{
321	let mut parachain_config = prepare_node_config(parachain_config);
322
323	let params = new_partial(&mut parachain_config, proof_recording_during_import)?;
324
325	let transaction_pool = params.transaction_pool.clone();
326	let mut task_manager = params.task_manager;
327
328	let client = params.client.clone();
329	let backend = params.backend.clone();
330
331	let block_import = params.other.0;
332	let slot_based_handle = params.other.1;
333	let relay_chain_interface = build_relay_chain_interface(
334		relay_chain_config,
335		parachain_config.prometheus_registry(),
336		collator_key.clone(),
337		collator_options.clone(),
338		&mut task_manager,
339	)
340	.await
341	.map_err(|e| sc_service::Error::Application(Box::new(e) as Box<_>))?;
342
343	let import_queue_service = params.import_queue.service();
344	let prometheus_registry = parachain_config.prometheus_registry().cloned();
345	let net_config = FullNetworkConfiguration::<Block, Hash, Net>::new(
346		&parachain_config.network,
347		prometheus_registry.clone(),
348	);
349
350	let best_hash = client.chain_info().best_hash;
351	let para_id = client
352		.runtime_api()
353		.parachain_id(best_hash)
354		.map_err(|e| sc_service::Error::Application(Box::new(e) as Box<_>))?;
355	tracing::info!("Parachain id: {:?}", para_id);
356
357	let (network, system_rpc_tx, tx_handler_controller, sync_service) =
358		build_network(BuildNetworkParams {
359			parachain_config: &parachain_config,
360			net_config,
361			client: client.clone(),
362			transaction_pool: transaction_pool.clone(),
363			para_id,
364			spawn_handle: task_manager.spawn_handle(),
365			relay_chain_interface: relay_chain_interface.clone(),
366			import_queue: params.import_queue,
367			metrics: Net::register_notification_metrics(
368				parachain_config.prometheus_config.as_ref().map(|config| &config.registry),
369			),
370			sybil_resistance_level: CollatorSybilResistance::Resistant,
371		})
372		.await?;
373
374	let keystore = params.keystore_container.keystore();
375	let rpc_builder = {
376		let client = client.clone();
377		Box::new(move |_| rpc_ext_builder(client.clone()))
378	};
379
380	let rpc_handlers = sc_service::spawn_tasks(sc_service::SpawnTasksParams {
381		rpc_builder,
382		client: client.clone(),
383		transaction_pool: transaction_pool.clone(),
384		task_manager: &mut task_manager,
385		config: parachain_config,
386		keystore: keystore.clone(),
387		backend: backend.clone(),
388		network: network.clone(),
389		sync_service: sync_service.clone(),
390		system_rpc_tx,
391		tx_handler_controller,
392		telemetry: None,
393		tracing_execute_block: Some(Arc::new(ParachainTracingExecuteBlock::new(client.clone()))),
394	})?;
395
396	let announce_block = {
397		let sync_service = sync_service.clone();
398		Arc::new(move |hash, data| sync_service.announce_block(hash, data))
399	};
400
401	let announce_block = wrap_announce_block
402		.map(|w| (w)(announce_block.clone()))
403		.unwrap_or_else(|| announce_block);
404
405	let overseer_handle = relay_chain_interface
406		.overseer_handle()
407		.map_err(|e| sc_service::Error::Application(Box::new(e)))?;
408
409	let recovery_handle: Box<dyn RecoveryHandle> = if fail_pov_recovery {
410		Box::new(FailingRecoveryHandle::new(overseer_handle.clone()))
411	} else {
412		Box::new(overseer_handle.clone())
413	};
414	let relay_chain_slot_duration = Duration::from_secs(6);
415
416	start_relay_chain_tasks(StartRelayChainTasksParams {
417		client: client.clone(),
418		announce_block: announce_block.clone(),
419		para_id,
420		relay_chain_interface: relay_chain_interface.clone(),
421		task_manager: &mut task_manager,
422		// Increase speed of recovery for testing purposes.
423		da_recovery_profile: DARecoveryProfile::Other(RecoveryDelayRange {
424			min: Duration::from_secs(1),
425			max: Duration::from_secs(5),
426		}),
427		import_queue: import_queue_service,
428		relay_chain_slot_duration,
429		recovery_handle,
430		sync_service: sync_service.clone(),
431		prometheus_registry: None,
432	})?;
433
434	let collator_peer_id = network.local_peer_id();
435	if let Some(collator_key) = collator_key {
436		let proposer = sc_basic_authorship::ProposerFactory::new(
437			task_manager.spawn_handle(),
438			client.clone(),
439			transaction_pool.clone(),
440			prometheus_registry.as_ref(),
441			None,
442		);
443
444		let collator_service = CollatorService::new(
445			client.clone(),
446			Arc::new(task_manager.spawn_handle()),
447			announce_block,
448			client.clone(),
449		);
450
451		let client_for_aura = client.clone();
452
453		if use_slot_based_collator {
454			tracing::info!(target: LOG_TARGET, "Starting block authoring with slot based authoring.");
455			let params = SlotBasedParams {
456				create_inherent_data_providers: move |_, ()| async move { Ok(()) },
457				block_import,
458				para_client: client.clone(),
459				para_backend: backend.clone(),
460				relay_client: relay_chain_interface,
461				code_hash_provider: move |block_hash| {
462					client_for_aura.code_at(block_hash).ok().map(|c| ValidationCode::from(c).hash())
463				},
464				keystore,
465				collator_key,
466				relay_chain_slot_duration,
467				para_id,
468				proposer,
469				collator_service,
470				authoring_duration: Duration::from_millis(2000),
471				reinitialize: false,
472				slot_offset: Duration::from_secs(1),
473				block_import_handle: slot_based_handle,
474				spawner: task_manager.spawn_essential_handle(),
475				export_pov: None,
476				max_pov_percentage: None,
477				collator_peer_id,
478			};
479
480			slot_based::run::<Block, AuthorityPair, _, _, _, _, _, _, _, _, _>(params);
481		} else {
482			tracing::info!(target: LOG_TARGET, "Starting block authoring with lookahead collator.");
483			let params = AuraParams {
484				create_inherent_data_providers: move |_, ()| async move { Ok(()) },
485				block_import,
486				para_client: client.clone(),
487				para_backend: backend.clone(),
488				relay_client: relay_chain_interface,
489				code_hash_provider: move |block_hash| {
490					client_for_aura.code_at(block_hash).ok().map(|c| ValidationCode::from(c).hash())
491				},
492				keystore,
493				collator_key,
494				collator_peer_id,
495				para_id,
496				overseer_handle,
497				relay_chain_slot_duration,
498				proposer,
499				collator_service,
500				authoring_duration: Duration::from_millis(2000),
501				reinitialize: false,
502				max_pov_percentage: None,
503			};
504
505			let fut = aura::run::<Block, AuthorityPair, _, _, _, _, _, _, _, _>(params);
506			task_manager.spawn_essential_handle().spawn("aura", None, fut);
507		}
508	}
509
510	Ok((task_manager, client, network, rpc_handlers, transaction_pool, backend))
511}
512
513/// A Cumulus test node instance used for testing.
514pub struct TestNode {
515	/// TaskManager's instance.
516	pub task_manager: TaskManager,
517	/// Client's instance.
518	pub client: Arc<Client>,
519	/// Node's network.
520	pub network: Arc<dyn NetworkService>,
521	/// The `MultiaddrWithPeerId` to this node. This is useful if you want to pass it as "boot
522	/// node" to other nodes.
523	pub addr: MultiaddrWithPeerId,
524	/// RPCHandlers to make RPC queries.
525	pub rpc_handlers: RpcHandlers,
526	/// Node's transaction pool
527	pub transaction_pool: TransactionPool,
528	/// Node's backend
529	pub backend: Arc<Backend>,
530}
531
532/// A builder to create a [`TestNode`].
533pub struct TestNodeBuilder {
534	para_id: ParaId,
535	tokio_handle: tokio::runtime::Handle,
536	key: Sr25519Keyring,
537	collator_key: Option<CollatorPair>,
538	parachain_nodes: Vec<MultiaddrWithPeerId>,
539	parachain_nodes_exclusive: bool,
540	relay_chain_nodes: Vec<MultiaddrWithPeerId>,
541	wrap_announce_block: Option<Box<dyn FnOnce(AnnounceBlockFn) -> AnnounceBlockFn>>,
542	storage_update_func_parachain: Option<Box<dyn Fn()>>,
543	storage_update_func_relay_chain: Option<Box<dyn Fn()>>,
544	relay_chain_mode: RelayChainMode,
545	endowed_accounts: Vec<AccountId>,
546	record_proof_during_import: bool,
547}
548
549impl TestNodeBuilder {
550	/// Create a new instance of `Self`.
551	///
552	/// `para_id` - The parachain id this node is running for.
553	/// `tokio_handle` - The tokio handler to use.
554	/// `key` - The key that will be used to generate the name and that will be passed as
555	/// `dev_seed`.
556	pub fn new(para_id: ParaId, tokio_handle: tokio::runtime::Handle, key: Sr25519Keyring) -> Self {
557		TestNodeBuilder {
558			key,
559			para_id,
560			tokio_handle,
561			collator_key: None,
562			parachain_nodes: Vec::new(),
563			parachain_nodes_exclusive: false,
564			relay_chain_nodes: Vec::new(),
565			wrap_announce_block: None,
566			storage_update_func_parachain: None,
567			storage_update_func_relay_chain: None,
568			endowed_accounts: Default::default(),
569			relay_chain_mode: RelayChainMode::Embedded,
570			record_proof_during_import: true,
571		}
572	}
573
574	/// Enable collator for this node.
575	pub fn enable_collator(mut self) -> Self {
576		let collator_key = CollatorPair::generate().0;
577		self.collator_key = Some(collator_key);
578		self
579	}
580
581	/// Instruct the node to exclusively connect to registered parachain nodes.
582	///
583	/// Parachain nodes can be registered using [`Self::connect_to_parachain_node`] and
584	/// [`Self::connect_to_parachain_nodes`].
585	pub fn exclusively_connect_to_registered_parachain_nodes(mut self) -> Self {
586		self.parachain_nodes_exclusive = true;
587		self
588	}
589
590	/// Make the node connect to the given parachain node.
591	///
592	/// By default the node will not be connected to any node or will be able to discover any other
593	/// node.
594	pub fn connect_to_parachain_node(mut self, node: &TestNode) -> Self {
595		self.parachain_nodes.push(node.addr.clone());
596		self
597	}
598
599	/// Make the node connect to the given parachain nodes.
600	///
601	/// By default the node will not be connected to any node or will be able to discover any other
602	/// node.
603	pub fn connect_to_parachain_nodes<'a>(
604		mut self,
605		nodes: impl IntoIterator<Item = &'a TestNode>,
606	) -> Self {
607		self.parachain_nodes.extend(nodes.into_iter().map(|n| n.addr.clone()));
608		self
609	}
610
611	/// Make the node connect to the given relay chain node.
612	///
613	/// By default the node will not be connected to any node or will be able to discover any other
614	/// node.
615	pub fn connect_to_relay_chain_node(
616		mut self,
617		node: &polkadot_test_service::PolkadotTestNode,
618	) -> Self {
619		self.relay_chain_nodes.push(node.addr.clone());
620		self
621	}
622
623	/// Make the node connect to the given relay chain nodes.
624	///
625	/// By default the node will not be connected to any node or will be able to discover any other
626	/// node.
627	pub fn connect_to_relay_chain_nodes<'a>(
628		mut self,
629		nodes: impl IntoIterator<Item = &'a polkadot_test_service::PolkadotTestNode>,
630	) -> Self {
631		self.relay_chain_nodes.extend(nodes.into_iter().map(|n| n.addr.clone()));
632		self
633	}
634
635	/// Wrap the announce block function of this node.
636	pub fn wrap_announce_block(
637		mut self,
638		wrap: impl FnOnce(AnnounceBlockFn) -> AnnounceBlockFn + 'static,
639	) -> Self {
640		self.wrap_announce_block = Some(Box::new(wrap));
641		self
642	}
643
644	/// Allows accessing the parachain storage before the test node is built.
645	pub fn update_storage_parachain(mut self, updater: impl Fn() + 'static) -> Self {
646		self.storage_update_func_parachain = Some(Box::new(updater));
647		self
648	}
649
650	/// Allows accessing the relay chain storage before the test node is built.
651	pub fn update_storage_relay_chain(mut self, updater: impl Fn() + 'static) -> Self {
652		self.storage_update_func_relay_chain = Some(Box::new(updater));
653		self
654	}
655
656	/// Connect to full node via RPC.
657	pub fn use_external_relay_chain_node_at_url(mut self, network_address: Url) -> Self {
658		self.relay_chain_mode = RelayChainMode::ExternalRpc(vec![network_address]);
659		self
660	}
661
662	/// Connect to full node via RPC.
663	pub fn use_external_relay_chain_node_at_port(mut self, port: u16) -> Self {
664		let mut localhost_url =
665			Url::parse("ws://localhost").expect("Should be able to parse localhost Url");
666		localhost_url.set_port(Some(port)).expect("Should be able to set port");
667		self.relay_chain_mode = RelayChainMode::ExternalRpc(vec![localhost_url]);
668		self
669	}
670
671	/// Accounts which will have an initial balance.
672	pub fn endowed_accounts(mut self, accounts: Vec<AccountId>) -> TestNodeBuilder {
673		self.endowed_accounts = accounts;
674		self
675	}
676
677	/// Record proofs during import.
678	pub fn import_proof_recording(mut self, should_record_proof: bool) -> TestNodeBuilder {
679		self.record_proof_during_import = should_record_proof;
680		self
681	}
682
683	/// Build the [`TestNode`].
684	pub async fn build(self) -> TestNode {
685		let parachain_config = node_config(
686			self.storage_update_func_parachain.unwrap_or_else(|| Box::new(|| ())),
687			self.tokio_handle.clone(),
688			self.key,
689			self.parachain_nodes,
690			self.parachain_nodes_exclusive,
691			self.para_id,
692			self.collator_key.is_some(),
693			self.endowed_accounts,
694		)
695		.expect("could not generate Configuration");
696
697		let mut relay_chain_config = polkadot_test_service::node_config(
698			self.storage_update_func_relay_chain.unwrap_or_else(|| Box::new(|| ())),
699			self.tokio_handle,
700			self.key,
701			self.relay_chain_nodes,
702			false,
703		);
704
705		let collator_options = CollatorOptions {
706			relay_chain_mode: self.relay_chain_mode,
707			embedded_dht_bootnode: true,
708			dht_bootnode_discovery: true,
709		};
710
711		relay_chain_config.network.node_name =
712			format!("{} (relay chain)", relay_chain_config.network.node_name);
713
714		let (task_manager, client, network, rpc_handlers, transaction_pool, backend) =
715			match relay_chain_config.network.network_backend {
716				sc_network::config::NetworkBackendType::Libp2p =>
717					start_node_impl::<_, sc_network::NetworkWorker<_, _>>(
718						parachain_config,
719						self.collator_key,
720						relay_chain_config,
721						self.wrap_announce_block,
722						false,
723						|_| Ok(jsonrpsee::RpcModule::new(())),
724						collator_options,
725						self.record_proof_during_import,
726						false,
727					)
728					.await
729					.expect("could not create Cumulus test service"),
730				sc_network::config::NetworkBackendType::Litep2p =>
731					start_node_impl::<_, sc_network::Litep2pNetworkBackend>(
732						parachain_config,
733						self.collator_key,
734						relay_chain_config,
735						self.wrap_announce_block,
736						false,
737						|_| Ok(jsonrpsee::RpcModule::new(())),
738						collator_options,
739						self.record_proof_during_import,
740						false,
741					)
742					.await
743					.expect("could not create Cumulus test service"),
744			};
745		let peer_id = network.local_peer_id();
746		let multiaddr = polkadot_test_service::get_listen_address(network.clone()).await;
747		let addr = MultiaddrWithPeerId { multiaddr, peer_id };
748
749		TestNode { task_manager, client, network, addr, rpc_handlers, transaction_pool, backend }
750	}
751}
752
753/// Create a Cumulus `Configuration`.
754///
755/// By default a TCP socket will be used, therefore you need to provide nodes if you want the
756/// node to be connected to other nodes.
757///
758/// If `nodes_exclusive` is `true`, the node will only connect to the given `nodes` and not to any
759/// other node.
760///
761/// The `storage_update_func` can be used to make adjustments to the runtime genesis.
762pub fn node_config(
763	storage_update_func: impl Fn(),
764	tokio_handle: tokio::runtime::Handle,
765	key: Sr25519Keyring,
766	nodes: Vec<MultiaddrWithPeerId>,
767	nodes_exclusive: bool,
768	para_id: ParaId,
769	is_collator: bool,
770	endowed_accounts: Vec<AccountId>,
771) -> Result<Configuration, ServiceError> {
772	let base_path = BasePath::new_temp_dir()?;
773	let root = base_path.path().join(format!("cumulus_test_service_{}", key));
774	let role = if is_collator { Role::Authority } else { Role::Full };
775	let key_seed = key.to_seed();
776	let mut spec = Box::new(chain_spec::get_chain_spec_with_extra_endowed(
777		Some(para_id),
778		endowed_accounts,
779		cumulus_test_runtime::WASM_BINARY.expect("WASM binary was not built, please build it!"),
780	));
781
782	let mut storage = spec.as_storage_builder().build_storage().expect("could not build storage");
783
784	BasicExternalities::execute_with_storage(&mut storage, storage_update_func);
785	spec.set_storage(storage);
786
787	let mut network_config = NetworkConfiguration::new(
788		format!("{} (parachain)", key_seed),
789		"network/test/0.1",
790		Default::default(),
791		None,
792	);
793
794	if nodes_exclusive {
795		network_config.default_peers_set.reserved_nodes = nodes;
796		network_config.default_peers_set.non_reserved_mode =
797			sc_network::config::NonReservedPeerMode::Deny;
798	} else {
799		network_config.boot_nodes = nodes;
800	}
801
802	network_config.allow_non_globals_in_dht = true;
803
804	let addr: multiaddr::Multiaddr = "/ip4/127.0.0.1/tcp/0".parse().expect("valid address; qed");
805	network_config.listen_addresses.push(addr.clone());
806	network_config.transport =
807		TransportConfig::Normal { enable_mdns: false, allow_private_ip: true };
808
809	Ok(Configuration {
810		impl_name: "cumulus-test-node".to_string(),
811		impl_version: "0.1".to_string(),
812		role,
813		tokio_handle,
814		transaction_pool: Default::default(),
815		network: network_config,
816		keystore: KeystoreConfig::InMemory,
817		database: DatabaseSource::RocksDb { path: root.join("db"), cache_size: 128 },
818		trie_cache_maximum_size: Some(64 * 1024 * 1024),
819		warm_up_trie_cache: None,
820		state_pruning: Some(PruningMode::ArchiveAll),
821		blocks_pruning: BlocksPruning::KeepAll,
822		chain_spec: spec,
823		executor: ExecutorConfiguration {
824			wasm_method: WasmExecutionMethod::Compiled {
825				instantiation_strategy:
826					sc_executor_wasmtime::InstantiationStrategy::PoolingCopyOnWrite,
827			},
828			..ExecutorConfiguration::default()
829		},
830		rpc: RpcConfiguration {
831			addr: None,
832			max_connections: Default::default(),
833			cors: None,
834			methods: Default::default(),
835			max_request_size: Default::default(),
836			max_response_size: Default::default(),
837			id_provider: None,
838			max_subs_per_conn: Default::default(),
839			port: 9945,
840			message_buffer_capacity: Default::default(),
841			batch_config: RpcBatchRequestConfig::Unlimited,
842			rate_limit: None,
843			rate_limit_whitelisted_ips: Default::default(),
844			rate_limit_trust_proxy_headers: Default::default(),
845			request_logger_limit: 1024,
846		},
847		prometheus_config: None,
848		telemetry_endpoints: None,
849		offchain_worker: OffchainWorkerConfig { enabled: true, indexing_enabled: false },
850		force_authoring: false,
851		disable_grandpa: false,
852		dev_key_seed: Some(key_seed),
853		tracing_targets: None,
854		tracing_receiver: Default::default(),
855		announce_block: true,
856		data_path: root,
857		base_path,
858		wasm_runtime_overrides: None,
859	})
860}
861
862impl TestNode {
863	/// Wait for `count` blocks to be imported in the node and then exit. This function will not
864	/// return if no blocks are ever created, thus you should restrict the maximum amount of time of
865	/// the test execution.
866	pub fn wait_for_blocks(&self, count: usize) -> impl Future<Output = ()> {
867		self.client.wait_for_blocks(count)
868	}
869
870	/// Send an extrinsic to this node.
871	pub async fn send_extrinsic(
872		&self,
873		function: impl Into<runtime::RuntimeCall>,
874		caller: Sr25519Keyring,
875	) -> Result<RpcTransactionOutput, RpcTransactionError> {
876		let extrinsic = construct_extrinsic(&self.client, function, caller.pair(), Some(0));
877
878		self.rpc_handlers.send_transaction(extrinsic.into()).await
879	}
880
881	/// Register a parachain at this relay chain.
882	pub async fn schedule_upgrade(&self, validation: Vec<u8>) -> Result<(), RpcTransactionError> {
883		let call = frame_system::Call::set_code { code: validation };
884
885		self.send_extrinsic(
886			runtime::SudoCall::sudo_unchecked_weight {
887				call: Box::new(call.into()),
888				weight: Weight::from_parts(1_000, 0),
889			},
890			Sr25519Keyring::Alice,
891		)
892		.await
893		.map(drop)
894	}
895}
896
897/// Fetch account nonce for key pair
898pub fn fetch_nonce(client: &Client, account: sp_core::sr25519::Public) -> u32 {
899	let best_hash = client.chain_info().best_hash;
900	client
901		.runtime_api()
902		.account_nonce(best_hash, account.into())
903		.expect("Fetching account nonce works; qed")
904}
905
906/// Construct an extrinsic that can be applied to the test runtime.
907pub fn construct_extrinsic(
908	client: &Client,
909	function: impl Into<runtime::RuntimeCall>,
910	caller: sp_core::sr25519::Pair,
911	nonce: Option<u32>,
912) -> runtime::UncheckedExtrinsic {
913	let function = function.into();
914	let current_block_hash = client.info().best_hash;
915	let current_block = client.info().best_number.saturated_into();
916	let genesis_block = client.hash(0).unwrap().unwrap();
917	let nonce = nonce.unwrap_or_else(|| fetch_nonce(client, caller.public()));
918	let period = runtime::BlockHashCount::get()
919		.checked_next_power_of_two()
920		.map(|c| c / 2)
921		.unwrap_or(2) as u64;
922	let tip = 0;
923	let tx_ext: runtime::TxExtension = (
924		frame_system::AuthorizeCall::<runtime::Runtime>::new(),
925		frame_system::CheckNonZeroSender::<runtime::Runtime>::new(),
926		frame_system::CheckSpecVersion::<runtime::Runtime>::new(),
927		frame_system::CheckGenesis::<runtime::Runtime>::new(),
928		frame_system::CheckEra::<runtime::Runtime>::from(generic::Era::mortal(
929			period,
930			current_block,
931		)),
932		frame_system::CheckNonce::<runtime::Runtime>::from(nonce),
933		frame_system::CheckWeight::<runtime::Runtime>::new(),
934		pallet_transaction_payment::ChargeTransactionPayment::<runtime::Runtime>::from(tip),
935	)
936		.into();
937	let raw_payload = runtime::SignedPayload::from_raw(
938		function.clone(),
939		tx_ext.clone(),
940		((), (), runtime::VERSION.spec_version, genesis_block, current_block_hash, (), (), ()),
941	);
942	let signature = raw_payload.using_encoded(|e| caller.sign(e));
943	runtime::UncheckedExtrinsic::new_signed(
944		function,
945		MultiAddress::Id(caller.public().into()),
946		runtime::Signature::Sr25519(signature),
947		tx_ext,
948	)
949}
950
951/// Run a relay-chain validator node.
952///
953/// This is essentially a wrapper around
954/// [`run_validator_node`](polkadot_test_service::run_validator_node).
955pub fn run_relay_chain_validator_node(
956	tokio_handle: tokio::runtime::Handle,
957	key: Sr25519Keyring,
958	storage_update_func: impl Fn(),
959	boot_nodes: Vec<MultiaddrWithPeerId>,
960	port: Option<u16>,
961) -> polkadot_test_service::PolkadotTestNode {
962	let mut config = polkadot_test_service::node_config(
963		storage_update_func,
964		tokio_handle.clone(),
965		key,
966		boot_nodes,
967		true,
968	);
969
970	if let Some(port) = port {
971		config.rpc.addr = Some(vec![RpcEndpoint {
972			batch_config: config.rpc.batch_config,
973			cors: config.rpc.cors.clone(),
974			listen_addr: SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, port)),
975			max_connections: config.rpc.max_connections,
976			max_payload_in_mb: config.rpc.max_request_size,
977			max_payload_out_mb: config.rpc.max_response_size,
978			max_subscriptions_per_connection: config.rpc.max_subs_per_conn,
979			max_buffer_capacity_per_connection: config.rpc.message_buffer_capacity,
980			rpc_methods: config.rpc.methods,
981			rate_limit: config.rpc.rate_limit,
982			rate_limit_trust_proxy_headers: config.rpc.rate_limit_trust_proxy_headers,
983			rate_limit_whitelisted_ips: config.rpc.rate_limit_whitelisted_ips.clone(),
984			retry_random_port: true,
985			is_optional: false,
986		}]);
987	}
988
989	let mut workers_path = std::env::current_exe().unwrap();
990	workers_path.pop();
991	workers_path.pop();
992
993	tokio_handle.block_on(async move {
994		polkadot_test_service::run_validator_node(config, Some(workers_path)).await
995	})
996}