Improve remote-externalities

utils/frame/remote-externalities/Cargo.toml

@@ -25,9 +25,10 @@ codec = { package = "parity-scale-codec", version = "2.0.0" }
 
 sp-io = { version = "3.0.0", path = "../../../primitives/io" }
 sp-core = { version = "3.0.0", path = "../../../primitives/core" }
+sp-runtime = { version = "3.0.0", path = "../../../primitives/runtime" }
 
 [dev-dependencies]
-async-std = { version = "1.6.5", features = ["attributes"] }
+tokio = { version = "0.2", features = ["macros"] }
 
 [features]
 remote-test = []

utils/frame/remote-externalities/src/lib.rs

@@ -106,7 +106,7 @@ use std::{
 	path::{Path, PathBuf},
 };
 use log::*;
-use sp_core::{hashing::twox_128};
+use sp_core::hashing::twox_128;
 pub use sp_io::TestExternalities;
 use sp_core::{
 	hexdisplay::HexDisplay,
@@ -115,27 +115,35 @@ use sp_core::{
 use codec::{Encode, Decode};
 use jsonrpsee_http_client::{HttpClient, HttpConfig};
 
+use sp_runtime::{
+	generic::BlockId,
+    traits::{
+		Block as BlockT, NumberFor,
+    }
+};
+
+// TODO: Make KeyPair generic
 type KeyPair = (StorageKey, StorageData);
-type Hash = sp_core::H256;
-// TODO: make these two generic.
 
 const LOG_TARGET: &str = "remote-ext";
 const TARGET: &str = "http://localhost:9933";
 
 jsonrpsee_proc_macros::rpc_client_api! {
-	RpcApi {
+	RpcApi<B: BlockT> {
 		#[rpc(method = "state_getPairs", positional_params)]
-		fn storage_pairs(prefix: StorageKey, hash: Option<Hash>) -> Vec<(StorageKey, StorageData)>;
+		fn storage_pairs(prefix: StorageKey, hash: Option<B::Hash>) -> Vec<(StorageKey, StorageData)>;
 		#[rpc(method = "chain_getFinalizedHead")]
-		fn finalized_head() -> Hash;
+		fn finalized_head() -> B::Hash;
+		#[rpc(method = "chain_getBlockHash")]
+		fn block_hash(number: NumberFor<B>) -> B::Hash;
 	}
 }
 
 /// The execution mode.
 #[derive(Clone)]
-pub enum Mode {
+pub enum Mode<B: BlockT> {
 	/// Online.
-	Online(OnlineConfig),
+	Online(OnlineConfig<B>),
 	/// Offline. Uses a cached file and needs not any client config.
 	Offline(OfflineConfig),
 }
@@ -153,24 +161,24 @@ pub struct OfflineConfig {
 ///
 /// A cache config may be present and will be written to in that case.
 #[derive(Clone)]
-pub struct OnlineConfig {
+pub struct OnlineConfig<B: BlockT> {
 	/// The HTTP uri to use.
 	pub uri: String,
 	/// The block number at which to connect. Will be latest finalized head if not provided.
-	pub at: Option<Hash>,
+	pub at: Option<BlockId<B>>,
 	/// An optional cache file to WRITE to, not for reading. Not cached if set to `None`.
 	pub cache: Option<CacheConfig>,
 	/// The modules to scrape. If empty, entire chain state will be scraped.
 	pub modules: Vec<String>,
 }
 
-impl Default for OnlineConfig {
+impl<B: BlockT> Default for OnlineConfig<B> {
 	fn default() -> Self {
 		Self { uri: TARGET.to_owned(), at: None, cache: None, modules: Default::default() }
 	}
 }
 
-impl OnlineConfig {
+impl<B: BlockT> OnlineConfig<B> {
 	/// Return a new http rpc client.
 	fn rpc(&self) -> HttpClient {
 		HttpClient::new(&self.uri, HttpConfig { max_request_body_size: u32::MAX })
@@ -202,30 +210,30 @@ impl CacheConfig {
 }
 
 /// Builder for remote-externalities.
-pub struct Builder {
+pub struct Builder<B: BlockT> {
 	inject: Vec<KeyPair>,
-	mode: Mode,
+	mode: Mode<B>,
 }
 
-impl Default for Builder {
+impl<B: BlockT> Default for Builder<B> {
 	fn default() -> Self {
 		Self {
 			inject: Default::default(),
-			mode: Mode::Online(OnlineConfig::default())
+			mode: Mode::Online(OnlineConfig::default()),
 		}
 	}
 }
 
 // Mode methods
-impl Builder {
-	fn as_online(&self) -> &OnlineConfig {
+impl<B: BlockT> Builder<B> {
+	fn as_online(&self) -> &OnlineConfig<B> {
 		match &self.mode {
 			Mode::Online(config) => &config,
 			_ => panic!("Unexpected mode: Online"),
 		}
 	}
 
-	fn as_online_mut(&mut self) -> &mut OnlineConfig {
+	fn as_online_mut(&mut self) -> &mut OnlineConfig<B> {
 		match &mut self.mode {
 			Mode::Online(config) => config,
 			_ => panic!("Unexpected mode: Online"),
@@ -234,10 +242,10 @@ impl Builder {
 }
 
 // RPC methods
-impl Builder {
-	async fn rpc_get_head(&self) -> Result<Hash, &'static str> {
+impl<B: BlockT> Builder<B> {
+	async fn rpc_get_head(&self) -> Result<B::Hash, &'static str> {
 		trace!(target: LOG_TARGET, "rpc: finalized_head");
-		RpcApi::finalized_head(&self.as_online().rpc()).await.map_err(|e| {
+		RpcApi::<B>::finalized_head(&self.as_online().rpc()).await.map_err(|e| {
 			error!("Error = {:?}", e);
 			"rpc finalized_head failed."
 			})
@@ -249,18 +257,30 @@ impl Builder {
 	async fn rpc_get_pairs(
 		&self,
 		prefix: StorageKey,
-		at: Hash,
+		at: B::Hash,
 	) -> Result<Vec<KeyPair>, &'static str> {
 		trace!(target: LOG_TARGET, "rpc: storage_pairs: {:?} / {:?}", prefix, at);
-		RpcApi::storage_pairs(&self.as_online().rpc(), prefix, Some(at)).await.map_err(|e| {
+		RpcApi::<B>::storage_pairs(&self.as_online().rpc(), prefix, Some(at)).await.map_err(|e| {
 			error!("Error = {:?}", e);
 			"rpc storage_pairs failed"
 			})
 	}
+
+	/// Relay the request to `chain_getBlockHash` rpc endpoint.
+	async fn rpc_get_hash(
+		&self,
+		number: NumberFor<B>,
+	) -> Result<B::Hash, &'static str> {
+		trace!(target: LOG_TARGET, "rpc: block_hash: {:?}", number);
+		RpcApi::<B>::block_hash(&self.as_online().rpc(), number).await.map_err(|e| {
+			error!("Error = {:?}", e);
+			"rpc block_hash failed"
+			})
+	}
 }
 
 // Internal methods
-impl Builder {
+impl<B: BlockT> Builder<B> {
 	/// Save the given data as cache.
 	fn save_cache(&self, data: &[KeyPair], path: &Path) -> Result<(), &'static str> {
 		info!(target: LOG_TARGET, "writing to cache file {:?}", path);
@@ -275,6 +295,13 @@ impl Builder {
 		Decode::decode(&mut &*bytes).map_err(|_| "decode failed")
 	}
 
+	async fn block_id_to_hash(&self, block_id: BlockId<B>) -> Result<B::Hash, &'static str> {
+		Ok(match block_id {
+			BlockId::Hash(hash) => hash,
+			BlockId::Number(number) => self.rpc_get_hash(number).await?,
+		})
+	}
+
 	/// Build `Self` from a network node denoted by `uri`.
 	async fn load_remote(&self) -> Result<Vec<KeyPair>, &'static str> {
 		let config = self.as_online();
@@ -289,7 +316,7 @@ impl Builder {
 			let mut filtered_kv = vec![];
 			for f in config.modules.iter() {
 				let hashed_prefix = StorageKey(twox_128(f.as_bytes()).to_vec());
-				let module_kv = self.rpc_get_pairs(hashed_prefix.clone(), at).await?;
+				let module_kv = self.rpc_get_pairs(hashed_prefix.clone(), self.block_id_to_hash(at).await?).await?;
 				info!(
 					target: LOG_TARGET,
 					"downloaded data for module {} (count: {} / prefix: {:?}).",
@@ -302,7 +329,7 @@ impl Builder {
 			filtered_kv
 		} else {
 			info!(target: LOG_TARGET, "downloading data for all modules.");
-			self.rpc_get_pairs(StorageKey(vec![]), at).await?.into_iter().collect::<Vec<_>>()
+			self.rpc_get_pairs(StorageKey(vec![]), self.block_id_to_hash(at).await?).await?.into_iter().collect::<Vec<_>>()
 		};
 
 		Ok(keys_and_values)
@@ -312,7 +339,7 @@ impl Builder {
 		info!(target: LOG_TARGET, "initializing remote client to {:?}", self.as_online().uri);
 		if self.as_online().at.is_none() {
 			let at = self.rpc_get_head().await?;
-			self.as_online_mut().at = Some(at);
+			self.as_online_mut().at = Some(BlockId::Hash(at));
 		}
 		Ok(())
 	}
@@ -341,7 +368,7 @@ impl Builder {
 }
 
 // Public methods
-impl Builder {
+impl<B: BlockT> Builder<B> {
 	/// Create a new builder.
 	pub fn new() -> Self {
 		Default::default()
@@ -356,7 +383,7 @@ impl Builder {
 	}
 
 	/// Configure a cache to be used.
-	pub fn mode(mut self, mode: Mode) -> Self {
+	pub fn mode(mut self, mode: Mode<B>) -> Self {
 		self.mode = mode;
 		self
 	}
@@ -379,6 +406,9 @@ impl Builder {
 #[cfg(test)]
 mod tests {
 	use super::*;
+	pub type Header = sp_runtime::generic::Header<u32, sp_runtime::traits::BlakeTwo256>;
+	pub type Block = sp_runtime::generic::Block<Header, UncheckedExtrinsic>;
+	pub type UncheckedExtrinsic = sp_runtime::generic::UncheckedExtrinsic<u32, (), (), ()>;
 
 	fn init_logger() {
 		let _ = env_logger::Builder::from_default_env()
@@ -387,10 +417,10 @@ mod tests {
 			.try_init();
 	}
 
-	#[async_std::test]
+	#[tokio::test]
 	async fn can_build_one_pallet() {
 		init_logger();
-		Builder::new()
+		Builder::<Block>::new()
 			.mode(Mode::Online(OnlineConfig {
 				modules: vec!["Proxy".into()],
 				..Default::default()
@@ -401,10 +431,10 @@ mod tests {
 			.execute_with(|| {});
 	}
 
-	#[async_std::test]
+	#[tokio::test]
 	async fn can_load_cache() {
 		init_logger();
-		Builder::new()
+		Builder::<Block>::new()
 			.mode(Mode::Offline(OfflineConfig {
 				cache: CacheConfig { name: "proxy_test".into(), ..Default::default() },
 			}))
@@ -414,10 +444,10 @@ mod tests {
 			.execute_with(|| {});
 	}
 
-	#[async_std::test]
+	#[tokio::test]
 	async fn can_create_cache() {
 		init_logger();
-		Builder::new()
+		Builder::<Block>::new()
 			.mode(Mode::Online(OnlineConfig {
 				cache: Some(CacheConfig {
 					name: "test_cache_to_remove.bin".into(),
@@ -444,9 +474,9 @@ mod tests {
 		}
 	}
 
-	#[async_std::test]
+	#[tokio::test]
 	async fn can_build_all() {
 		init_logger();
-		Builder::new().build().await.unwrap().execute_with(|| {});
+		Builder::<Block>::new().build().await.unwrap().execute_with(|| {});
 	}
 }