From f54614e2562fad6d244392e9b7873a170193f923 Mon Sep 17 00:00:00 2001
From: Cecile Tonglet <cecile@parity.io>
Date: Fri, 3 Jul 2020 11:49:42 +0200
Subject: [PATCH] New testing helpers (#6555)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* Initial commit

Forked at: 424d5c722d72caa1235652dea2f0a24ef4d054c0
Parent branch: origin/master

* Add send_transaction to RpcHandlers

* Extension trait for RpcHandlers

* Revert "Add send_transaction to RpcHandlers"

This reverts commit 03c89e13d404bae3f3123387dd50f026061bca82.

* Add an extension trait for BlockchainEvents

* Update test-utils/client/src/lib.rs

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

* Update test-utils/client/src/lib.rs

* fix

* deps fix

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>
---
 substrate/Cargo.lock                   |  2 +
 substrate/test-utils/client/Cargo.toml |  2 +
 substrate/test-utils/client/src/lib.rs | 87 +++++++++++++++++++++++++-
 3 files changed, 88 insertions(+), 3 deletions(-)

diff --git a/substrate/Cargo.lock b/substrate/Cargo.lock
index 9ea9b89057..b5bbcd6954 100644
--- a/substrate/Cargo.lock
+++ b/substrate/Cargo.lock
@@ -8309,8 +8309,10 @@ dependencies = [
 name = "substrate-test-client"
 version = "2.0.0-rc4"
 dependencies = [
+ "futures 0.1.29",
  "futures 0.3.5",
  "hash-db",
+ "hex",
  "parity-scale-codec",
  "sc-client-api",
  "sc-client-db",
diff --git a/substrate/test-utils/client/Cargo.toml b/substrate/test-utils/client/Cargo.toml
index a9d8590f02..e9036bc77a 100644
--- a/substrate/test-utils/client/Cargo.toml
+++ b/substrate/test-utils/client/Cargo.toml
@@ -20,7 +20,9 @@ sc-executor = { version = "0.8.0-rc4", path = "../../client/executor" }
 sc-consensus = { version = "0.8.0-rc4", path = "../../client/consensus/common" }
 sc-service = { version = "0.8.0-rc4", default-features = false, features = ["test-helpers"],  path = "../../client/service" }
 futures = "0.3.4"
+futures01 = { package = "futures", version = "0.1.29" }
 hash-db = "0.15.2"
+hex = "0.4"
 sp-keyring = { version = "2.0.0-rc4", path = "../../primitives/keyring" }
 codec = { package = "parity-scale-codec", version = "1.3.1" }
 sp-core = { version = "2.0.0-rc4", path = "../../primitives/core" }
diff --git a/substrate/test-utils/client/src/lib.rs b/substrate/test-utils/client/src/lib.rs
index 2ab9e4066d..fef9acd9d2 100644
--- a/substrate/test-utils/client/src/lib.rs
+++ b/substrate/test-utils/client/src/lib.rs
@@ -36,14 +36,17 @@ pub use sp_keyring::{
 pub use sp_core::{traits::BareCryptoStorePtr, tasks::executor as tasks_executor};
 pub use sp_runtime::{Storage, StorageChild};
 pub use sp_state_machine::ExecutionStrategy;
-pub use sc_service::client;
+pub use sc_service::{RpcHandlers, RpcSession, client};
 pub use self::client_ext::{ClientExt, ClientBlockImportExt};
 
+use std::pin::Pin;
 use std::sync::Arc;
-use std::collections::HashMap;
+use std::collections::{HashSet, HashMap};
+use futures::{future::{Future, FutureExt}, stream::StreamExt};
 use sp_core::storage::ChildInfo;
-use sp_runtime::traits::{Block as BlockT, BlakeTwo256};
+use sp_runtime::{OpaqueExtrinsic, codec::Encode, traits::{Block as BlockT, BlakeTwo256}};
 use sc_service::client::{LocalCallExecutor, ClientConfig};
+use sc_client_api::BlockchainEvents;
 
 /// Test client light database backend.
 pub type LightBackend<Block> = sc_light::Backend<
@@ -255,3 +258,81 @@ impl<Block: BlockT, E, Backend, G: GenesisInit> TestClientBuilder<
 		self.build_with_executor(executor)
 	}
 }
+
+/// An extension trait for `RpcHandlers`.
+pub trait RpcHandlersExt {
+	/// Send a transaction through the RpcHandlers.
+	fn send_transaction(
+		&self,
+		extrinsic: OpaqueExtrinsic,
+	) -> Pin<Box<dyn Future<
+		Output = (
+			Option<String>,
+			RpcSession,
+			futures01::sync::mpsc::Receiver<String>,
+		),
+	> + Send>>;
+}
+
+impl RpcHandlersExt for RpcHandlers {
+	fn send_transaction(
+		&self,
+		extrinsic: OpaqueExtrinsic,
+	) -> Pin<Box<dyn Future<
+		Output = (
+			Option<String>,
+			RpcSession,
+			futures01::sync::mpsc::Receiver<String>,
+		),
+	> + Send>> {
+		let (tx, rx) = futures01::sync::mpsc::channel(0);
+		let mem = RpcSession::new(tx.into());
+		Box::pin(self
+			.rpc_query(
+				&mem,
+				&format!(
+					r#"{{
+						"jsonrpc": "2.0",
+						"method": "author_submitExtrinsic",
+						"params": ["0x{}"],
+						"id": 0
+					}}"#,
+					hex::encode(extrinsic.encode())
+				),
+			)
+			.map(move |res| (res, mem, rx)))
+	}
+}
+
+/// An extension trait for `BlockchainEvents`.
+pub trait BlockchainEventsExt<C, B>
+where
+	C: BlockchainEvents<B>,
+	B: BlockT,
+{
+	/// 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.
+	fn wait_for_blocks(&self, count: usize) -> Pin<Box<dyn Future<Output = ()> + Send>>;
+}
+
+impl<C, B> BlockchainEventsExt<C, B> for C
+where
+	C: BlockchainEvents<B>,
+	B: BlockT,
+{
+	fn wait_for_blocks(&self, count: usize) -> Pin<Box<dyn Future<Output = ()> + Send>> {
+		assert!(count > 0, "'count' argument must be greater than 0");
+
+		let mut import_notification_stream = self.import_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;
+				}
+			}
+		})
+	}
+}