jsonrpsee integration (#8783)

* Add tokio * No need to map CallError to CallError * jsonrpsee proc macros (#9673) * port error types to `JsonRpseeError` * migrate chain module to proc macro api * make it compile with proc macros * update branch * update branch * update to jsonrpsee master * port system rpc * port state rpc * port childstate & offchain * frame system rpc * frame transaction payment * bring back CORS hack to work with polkadot UI * port babe rpc * port manual seal rpc * port frame mmr rpc * port frame contracts rpc * port finality grandpa rpc * port sync state rpc * resolve a few TODO + no jsonrpc deps * Update bin/node/rpc-client/src/main.rs * Update bin/node/rpc-client/src/main.rs * Update bin/node/rpc-client/src/main.rs * Update bin/node/rpc-client/src/main.rs * Port over system_ rpc tests * Make it compile * Use prost 0.8 * Use prost 0.8 * Make it compile * Ignore more failing tests * Comment out WIP tests * fix nit in frame system api * Update lockfile * No more juggling tokio versions * No more wait_for_stop ? * Remove browser-testing * Arguments must be arrays * Use same argument names * Resolve todo: no wait_for_stop for WS server Add todo: is parse_rpc_result used? Cleanup imports * fmt * log * One test passes * update jsonrpsee * update jsonrpsee * cleanup rpc-servers crate * jsonrpsee: add host and origin filtering (#9787) * add access control in the jsonrpsee servers * use master * fix nits * rpc runtime_version safe * fix nits * fix grumbles * remove unused files * resolve some todos * jsonrpsee more cleanup (#9803) * more cleanup * resolve TODOs * fix some unwraps * remove type hints * update jsonrpsee * downgrade zeroize * pin jsonrpsee rev * remove unwrap nit * Comment out more tests that aren't ported * Comment out more tests * Fix tests after merge * Subscription test * Invalid nonce test * Pending exts * WIP removeExtrinsic test * Test remove_extrinsic * Make state test: should_return_storage work * Uncomment/fix the other non-subscription related state tests * test: author_insertKey * test: author_rotateKeys * Get rest of state tests passing * asyncify a little more * Add todo to note #msg change * Crashing test for has_session_keys * Fix error conversion to avoid stack overflows Port author_hasSessionKeys test fmt * test author_hasKey * Add two missing tests Add a check on the return type Add todos for James's concerns * RPC tests for state, author and system (#9859) * Fix test runner * Impl Default for SubscriptionTaskExecutor * Keep the minimul amount of code needed to compile tests * Re-instate `RpcSession` (for now) * cleanup * Port over RPC tests * Add tokio * No need to map CallError to CallError * Port over system_ rpc tests * Make it compile * Use prost 0.8 * Use prost 0.8 * Make it compile * Ignore more failing tests * Comment out WIP tests * Update lockfile * No more juggling tokio versions * No more wait_for_stop ? * Remove browser-testing * Arguments must be arrays * Use same argument names * Resolve todo: no wait_for_stop for WS server Add todo: is parse_rpc_result used? Cleanup imports * fmt * log * One test passes * Comment out more tests that aren't ported * Comment out more tests * Fix tests after merge * Subscription test * Invalid nonce test * Pending exts * WIP removeExtrinsic test * Test remove_extrinsic * Make state test: should_return_storage work * Uncomment/fix the other non-subscription related state tests * test: author_insertKey * test: author_rotateKeys * Get rest of state tests passing * asyncify a little more * Add todo to note #msg change * Crashing test for has_session_keys * Fix error conversion to avoid stack overflows Port author_hasSessionKeys test fmt * test author_hasKey * Add two missing tests Add a check on the return type Add todos for James's concerns * offchain rpc tests * Address todos * fmt Co-authored-by: James Wilson <james@jsdw.me> * fix drop in state test * update jsonrpsee * fix ignored system test * fix chain tests * remove some boiler plate * Port BEEFY RPC (#9883) * Merge master * Port beefy RPC (ty @niklas!) * trivial changes left over from merge * Remove unused code * Update jsonrpsee * fix build * make tests compile again * beefy update jsonrpsee * fix: respect rpc methods policy * update cargo.lock * update jsonrpsee * update jsonrpsee * downgrade error logs * update jsonrpsee * Fix typo * remove unused file * Better name * Port Babe RPC tests * Put docs back * Resolve todo * Port tests for System RPCs * Resolve todo * fix build * Updated jsonrpsee to current master * fix: port finality grandpa rpc tests * Move .into() outside of the match * more review grumbles * jsonrpsee: add `rpc handlers` back (#10245) * add back RpcHandlers * cargo fmt * fix docs * fix grumble: remove needless alloc * resolve TODO * fmt * Fix typo * grumble: Use constants based on BASE_ERROR * grumble: DRY whitelisted listening addresses grumble: s/JSONRPC/JSON-RPC/ * cleanup * grumbles: Making readers aware of the possibility of gaps * review grumbles * grumbles * remove notes from niklasad1 * Update `jsonrpsee` * fix: jsonrpsee features * jsonrpsee: fallback to random port in case the specified port failed (#10304) * jsonrpsee: fallback to random port * better comment * Update client/rpc-servers/src/lib.rs Co-authored-by: Maciej Hirsz <1096222+maciejhirsz@users.noreply.github.com> * Update client/rpc-servers/src/lib.rs Co-authored-by: Maciej Hirsz <1096222+maciejhirsz@users.noreply.github.com> * address grumbles * cargo fmt * addrs already slice Co-authored-by: Maciej Hirsz <1096222+maciejhirsz@users.noreply.github.com> * Update jsonrpsee to 092081a0a2b8904c6ebd2cd99e16c7bc13ffc3ae * lockfile * update jsonrpsee * fix warning * Don't fetch jsonrpsee from crates * make tests compile again * fix rpc tests * remove unused deps * update tokio * fix rpc tests again * fix: test runner `HttpServerBuilder::builder` fails unless it's called within tokio runtime * cargo fmt * grumbles: fix subscription aliases * make clippy happy * update remaining subscriptions alias * cleanup * cleanup * fix chain subscription: less boiler plate (#10285) * fix chain subscription: less boiler plate * fix bad merge * cargo fmt * Switch to jsonrpsee 0.5 * fix build * add missing features * fix nit: remove needless Box::pin * Integrate jsonrpsee metrics (#10395) * draft metrics impl * Use latest api * Add missing file * Http server metrics * cleanup * bump jsonrpsee * Remove `ServerMetrics` and use a single middleware for both connection counting (aka sessions) and call metrics. * fix build * remove needless Arc::clone * Update to jsonrpsee 0.6 * lolz * fix metrics * Revert "lolz" This reverts commit eed6c6a56e78d8e307b4950f4c52a1c3a2322ba1. * fix: in-memory rpc support subscriptions * commit Cargo.lock * Update tests to 0.7 * fix TODOs * ws server: generate subscriptionIDs as Strings Some libraries seems to expect the subscription IDs to be Strings, let's not break this in this PR. * Increase timeout * Port over tests * cleanup * Using error codes from the spec * fix clippy * cargo fmt * update jsonrpsee * fix nits * fix: rpc_query * enable custom subid gen through spawn_tasks * remove unsed deps * unify tokio deps * Revert "enable custom subid gen through spawn_tasks" This reverts commit 5c5eb70328fe39d154fdb55c56e637b4548cf470. * fix bad merge of `test-utils` * fix more nits * downgrade wasm-instrument to 0.1.0 * [jsonrpsee]: enable custom RPC subscription ID generatation (#10731) * enable custom subid gen through spawn_tasks * fix nits * Update client/service/src/builder.rs Co-authored-by: David <dvdplm@gmail.com> * add Poc; needs jsonrpsee pr * update jsonrpsee * add re-exports * add docs Co-authored-by: David <dvdplm@gmail.com> * cargo fmt * fmt * port RPC-API dev * Remove unused file * fix nit: remove async trait * fix doc links * fix merge nit: remove jsonrpc deps * kill namespace on rpc apis * companion for jsonrpsee v0.10 (#11158) * companion for jsonrpsee v0.10 * update versions v0.10.0 * add some fixes * spelling * fix spaces Co-authored-by: Niklas Adolfsson <niklasadolfsson1@gmail.com> * send error before subs are closed * fix unsubscribe method names: chain * fix tests * jsonrpc server: print binded local address * grumbles: kill SubscriptionTaskExecutor * Update client/sync-state-rpc/src/lib.rs Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com> * Update client/rpc/src/chain/chain_full.rs Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com> * Update client/rpc/src/chain/chain_full.rs Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com> * sync-state-rpc: kill anyhow * no more anyhow * remove todo * jsonrpsee: fix bad params in subscriptions. (#11251) * update jsonrpsee * fix error responses * revert error codes * dont do weird stuff in drop impl * rpc servers: remove needless clone * Remove silly constants * chore: update jsonrpsee v0.12 * commit Cargo.lock * deps: downgrade git2 * feat: CLI flag max subscriptions per connection * metrics: use old logging format * fix: read WS address from substrate output (#11379) Co-authored-by: Niklas Adolfsson <niklasadolfsson1@gmail.com> Co-authored-by: James Wilson <james@jsdw.me> Co-authored-by: Maciej Hirsz <hello@maciej.codes> Co-authored-by: Maciej Hirsz <1096222+maciejhirsz@users.noreply.github.com> Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>
2026-06-13 07:01:05 +00:00 · 2022-05-10 10:52:19 +02:00
parent e45d53552d
commit 29c0c6a4a8
93 changed files with 3813 additions and 5094 deletions
@@ -23,19 +23,22 @@
 use parking_lot::RwLock;
 use std::sync::Arc;

+use sc_rpc::SubscriptionTaskExecutor;
 use sp_runtime::traits::Block as BlockT;

-use futures::{task::SpawnError, FutureExt, SinkExt, StreamExt, TryFutureExt};
-use jsonrpc_derive::rpc;
-use jsonrpc_pubsub::{manager::SubscriptionManager, typed::Subscriber, SubscriptionId};
+use futures::{task::SpawnError, FutureExt, StreamExt};
+use jsonrpsee::{
+	core::{async_trait, Error as JsonRpseeError, RpcResult},
+	proc_macros::rpc,
+	types::{error::CallError, ErrorObject},
+	PendingSubscription,
+};
 use log::warn;

 use beefy_gadget::notification::{BeefyBestBlockStream, BeefySignedCommitmentStream};

 mod notification;

-type FutureResult<T> = jsonrpc_core::BoxFuture<Result<T, jsonrpc_core::Error>>;
-
 #[derive(Debug, thiserror::Error)]
 /// Top-level error type for the RPC handler
 pub enum Error {
@@ -64,195 +67,149 @@ impl From<Error> for ErrorCode {
 	}
 }

-impl From<Error> for jsonrpc_core::Error {
+impl From<Error> for JsonRpseeError {
 	fn from(error: Error) -> Self {
-		let message = format!("{}", error);
+		let message = error.to_string();
 		let code = ErrorCode::from(error);
-		jsonrpc_core::Error {
+		JsonRpseeError::Call(CallError::Custom(ErrorObject::owned(
+			code as i32,
 			message,
-			code: jsonrpc_core::ErrorCode::ServerError(code as i64),
-			data: None,
-		}
+			None::<()>,
+		)))
 	}
 }

-/// Provides RPC methods for interacting with BEEFY.
-#[rpc]
+// Provides RPC methods for interacting with BEEFY.
+#[rpc(client, server)]
 pub trait BeefyApi<Notification, Hash> {
-	/// RPC Metadata
-	type Metadata;
-
 	/// Returns the block most recently finalized by BEEFY, alongside side its justification.
-	#[pubsub(
-		subscription = "beefy_justifications",
-		subscribe,
-		name = "beefy_subscribeJustifications"
+	#[subscription(
+		name = "beefy_subscribeJustifications" => "beefy_justifications",
+		unsubscribe = "beefy_unsubscribeJustifications",
+		item = Notification,
 	)]
-	fn subscribe_justifications(
-		&self,
-		metadata: Self::Metadata,
-		subscriber: Subscriber<Notification>,
-	);
-
-	/// Unsubscribe from receiving notifications about recently finalized blocks.
-	#[pubsub(
-		subscription = "beefy_justifications",
-		unsubscribe,
-		name = "beefy_unsubscribeJustifications"
-	)]
-	fn unsubscribe_justifications(
-		&self,
-		metadata: Option<Self::Metadata>,
-		id: SubscriptionId,
-	) -> jsonrpc_core::Result<bool>;
+	fn subscribe_justifications(&self);

 	/// Returns hash of the latest BEEFY finalized block as seen by this client.
 	///
 	/// The latest BEEFY block might not be available if the BEEFY gadget is not running
 	/// in the network or if the client is still initializing or syncing with the network.
 	/// In such case an error would be returned.
-	#[rpc(name = "beefy_getFinalizedHead")]
-	fn latest_finalized(&self) -> FutureResult<Hash>;
+	#[method(name = "beefy_getFinalizedHead")]
+	async fn latest_finalized(&self) -> RpcResult<Hash>;
 }

 /// Implements the BeefyApi RPC trait for interacting with BEEFY.
 pub struct BeefyRpcHandler<Block: BlockT> {
 	signed_commitment_stream: BeefySignedCommitmentStream<Block>,
 	beefy_best_block: Arc<RwLock<Option<Block::Hash>>>,
-	manager: SubscriptionManager,
+	executor: SubscriptionTaskExecutor,
 }

-impl<Block: BlockT> BeefyRpcHandler<Block> {
+impl<Block> BeefyRpcHandler<Block>
+where
+	Block: BlockT,
+{
 	/// Creates a new BeefyRpcHandler instance.
-	pub fn new<E>(
+	pub fn new(
 		signed_commitment_stream: BeefySignedCommitmentStream<Block>,
 		best_block_stream: BeefyBestBlockStream<Block>,
-		executor: E,
-	) -> Result<Self, Error>
-	where
-		E: futures::task::Spawn + Send + Sync + 'static,
-	{
+		executor: SubscriptionTaskExecutor,
+	) -> Result<Self, Error> {
 		let beefy_best_block = Arc::new(RwLock::new(None));

 		let stream = best_block_stream.subscribe();
 		let closure_clone = beefy_best_block.clone();
 		let future = stream.for_each(move |best_beefy| {
 			let async_clone = closure_clone.clone();
-			async move {
-				*async_clone.write() = Some(best_beefy);
-			}
+			async move { *async_clone.write() = Some(best_beefy) }
 		});

-		executor
-			.spawn_obj(futures::task::FutureObj::new(Box::pin(future)))
-			.map_err(|e| {
-				log::error!("Failed to spawn BEEFY RPC background task; err: {}", e);
-				e
-			})?;
-
-		let manager = SubscriptionManager::new(Arc::new(executor));
-		Ok(Self { signed_commitment_stream, beefy_best_block, manager })
+		executor.spawn("substrate-rpc-subscription", Some("rpc"), future.map(drop).boxed());
+		Ok(Self { signed_commitment_stream, beefy_best_block, executor })
 	}
 }

-impl<Block> BeefyApi<notification::EncodedSignedCommitment, Block::Hash> for BeefyRpcHandler<Block>
+#[async_trait]
+impl<Block> BeefyApiServer<notification::EncodedSignedCommitment, Block::Hash>
+	for BeefyRpcHandler<Block>
 where
 	Block: BlockT,
 {
-	type Metadata = sc_rpc::Metadata;
-
-	fn subscribe_justifications(
-		&self,
-		_metadata: Self::Metadata,
-		subscriber: Subscriber<notification::EncodedSignedCommitment>,
-	) {
+	fn subscribe_justifications(&self, pending: PendingSubscription) {
 		let stream = self
 			.signed_commitment_stream
 			.subscribe()
-			.map(|x| Ok::<_, ()>(Ok(notification::EncodedSignedCommitment::new::<Block>(x))));
+			.map(|sc| notification::EncodedSignedCommitment::new::<Block>(sc));

-		self.manager.add(subscriber, |sink| {
-			stream
-				.forward(sink.sink_map_err(|e| warn!("Error sending notifications: {:?}", e)))
-				.map(|_| ())
-		});
+		let fut = async move {
+			if let Some(mut sink) = pending.accept() {
+				sink.pipe_from_stream(stream).await;
+			}
+		}
+		.boxed();
+
+		self.executor
+			.spawn("substrate-rpc-subscription", Some("rpc"), fut.map(drop).boxed());
 	}

-	fn unsubscribe_justifications(
-		&self,
-		_metadata: Option<Self::Metadata>,
-		id: SubscriptionId,
-	) -> jsonrpc_core::Result<bool> {
-		Ok(self.manager.cancel(id))
-	}
-
-	fn latest_finalized(&self) -> FutureResult<Block::Hash> {
-		let result: Result<Block::Hash, jsonrpc_core::Error> = self
-			.beefy_best_block
+	async fn latest_finalized(&self) -> RpcResult<Block::Hash> {
+		self.beefy_best_block
 			.read()
 			.as_ref()
 			.cloned()
-			.ok_or_else(|| Error::EndpointNotReady.into());
-		let future = async move { result }.boxed();
-		future.map_err(jsonrpc_core::Error::from).boxed()
+			.ok_or(Error::EndpointNotReady)
+			.map_err(Into::into)
 	}
 }

 #[cfg(test)]
 mod tests {
 	use super::*;
-	use jsonrpc_core::{types::Params, Notification, Output};

-	use beefy_gadget::notification::{BeefySignedCommitment, BeefySignedCommitmentSender};
+	use beefy_gadget::notification::{
+		BeefyBestBlockStream, BeefySignedCommitment, BeefySignedCommitmentSender,
+	};
 	use beefy_primitives::{known_payload_ids, Payload};
 	use codec::{Decode, Encode};
+	use jsonrpsee::{types::EmptyParams, RpcModule};
 	use sp_runtime::traits::{BlakeTwo256, Hash};
 	use substrate_test_runtime_client::runtime::Block;

-	fn setup_io_handler(
-	) -> (jsonrpc_core::MetaIoHandler<sc_rpc::Metadata>, BeefySignedCommitmentSender<Block>) {
+	fn setup_io_handler() -> (RpcModule<BeefyRpcHandler<Block>>, BeefySignedCommitmentSender<Block>)
+	{
 		let (_, stream) = BeefyBestBlockStream::<Block>::channel();
 		setup_io_handler_with_best_block_stream(stream)
 	}

 	fn setup_io_handler_with_best_block_stream(
 		best_block_stream: BeefyBestBlockStream<Block>,
-	) -> (jsonrpc_core::MetaIoHandler<sc_rpc::Metadata>, BeefySignedCommitmentSender<Block>) {
+	) -> (RpcModule<BeefyRpcHandler<Block>>, BeefySignedCommitmentSender<Block>) {
 		let (commitment_sender, commitment_stream) =
 			BeefySignedCommitmentStream::<Block>::channel();

-		let handler: BeefyRpcHandler<Block> = BeefyRpcHandler::new(
+		let handler = BeefyRpcHandler::new(
 			commitment_stream,
 			best_block_stream,
-			sc_rpc::testing::TaskExecutor,
+			sc_rpc::testing::test_executor(),
 		)
-		.unwrap();
+		.expect("Setting up the BEEFY RPC handler works");

-		let mut io = jsonrpc_core::MetaIoHandler::default();
-		io.extend_with(BeefyApi::to_delegate(handler));
-
-		(io, commitment_sender)
+		(handler.into_rpc(), commitment_sender)
 	}

-	fn setup_session() -> (sc_rpc::Metadata, futures::channel::mpsc::UnboundedReceiver<String>) {
-		let (tx, rx) = futures::channel::mpsc::unbounded();
-		let meta = sc_rpc::Metadata::new(tx);
-		(meta, rx)
-	}
-
-	#[test]
-	fn uninitialized_rpc_handler() {
-		let (io, _) = setup_io_handler();
-
+	#[tokio::test]
+	async fn uninitialized_rpc_handler() {
+		let (rpc, _) = setup_io_handler();
 		let request = r#"{"jsonrpc":"2.0","method":"beefy_getFinalizedHead","params":[],"id":1}"#;
-		let response = r#"{"jsonrpc":"2.0","error":{"code":1,"message":"BEEFY RPC endpoint not ready"},"id":1}"#;
+		let expected_response = r#"{"jsonrpc":"2.0","error":{"code":1,"message":"BEEFY RPC endpoint not ready"},"id":1}"#.to_string();
+		let (result, _) = rpc.raw_json_request(&request).await.unwrap();

-		let meta = sc_rpc::Metadata::default();
-		assert_eq!(Some(response.into()), io.handle_request_sync(request, meta));
+		assert_eq!(expected_response, result,);
 	}

-	#[test]
-	fn latest_finalized_rpc() {
+	#[tokio::test]
+	async fn latest_finalized_rpc() {
 		let (sender, stream) = BeefyBestBlockStream::<Block>::channel();
 		let (io, _) = setup_io_handler_with_best_block_stream(stream);

@@ -266,83 +223,78 @@ mod tests {
 			\"jsonrpc\":\"2.0\",\
 			\"result\":\"0x2f0039e93a27221fcf657fb877a1d4f60307106113e885096cb44a461cd0afbf\",\
 			\"id\":1\
-		}";
+		}"
+		.to_string();
 		let not_ready = "{\
 			\"jsonrpc\":\"2.0\",\
 			\"error\":{\"code\":1,\"message\":\"BEEFY RPC endpoint not ready\"},\
 			\"id\":1\
-		}";
+		}"
+		.to_string();

 		let deadline = std::time::Instant::now() + std::time::Duration::from_secs(2);
 		while std::time::Instant::now() < deadline {
-			let meta = sc_rpc::Metadata::default();
-			let response = io.handle_request_sync(request, meta);
-			// Retry "not ready" responses.
-			if response != Some(not_ready.into()) {
-				assert_eq!(response, Some(expected.into()));
+			let (response, _) = io.raw_json_request(request).await.expect("RPC requests work");
+			if response != not_ready {
+				assert_eq!(response, expected);
 				// Success
 				return
 			}
-			std::thread::sleep(std::time::Duration::from_millis(50));
+			std::thread::sleep(std::time::Duration::from_millis(50))
 		}
+
 		panic!(
 			"Deadline reached while waiting for best BEEFY block to update. Perhaps the background task is broken?"
 		);
 	}

-	#[test]
-	fn subscribe_and_unsubscribe_to_justifications() {
-		let (io, _) = setup_io_handler();
-		let (meta, _) = setup_session();
+	#[tokio::test]
+	async fn subscribe_and_unsubscribe_to_justifications() {
+		let (rpc, _) = setup_io_handler();

-		// Subscribe
-		let sub_request =
-			r#"{"jsonrpc":"2.0","method":"beefy_subscribeJustifications","params":[],"id":1}"#;
-		let resp = io.handle_request_sync(sub_request, meta.clone());
-		let resp: Output = serde_json::from_str(&resp.unwrap()).unwrap();
+		// Subscribe call.
+		let sub = rpc
+			.subscribe("beefy_subscribeJustifications", EmptyParams::new())
+			.await
+			.unwrap();

-		let sub_id = match resp {
-			Output::Success(success) => success.result,
-			_ => panic!(),
-		};
+		let ser_id = serde_json::to_string(sub.subscription_id()).unwrap();

 		// Unsubscribe
 		let unsub_req = format!(
-			r#"{{"jsonrpc":"2.0","method":"beefy_unsubscribeJustifications","params":[{}],"id":1}}"#,
-			sub_id
-		);
-		assert_eq!(
-			io.handle_request_sync(&unsub_req, meta.clone()),
-			Some(r#"{"jsonrpc":"2.0","result":true,"id":1}"#.into()),
+			"{{\"jsonrpc\":\"2.0\",\"method\":\"beefy_unsubscribeJustifications\",\"params\":[{}],\"id\":1}}",
+			ser_id
 		);
+		let (response, _) = rpc.raw_json_request(&unsub_req).await.unwrap();
+
+		assert_eq!(response, r#"{"jsonrpc":"2.0","result":true,"id":1}"#);

 		// Unsubscribe again and fail
-		assert_eq!(
-			io.handle_request_sync(&unsub_req, meta),
-			Some(r#"{"jsonrpc":"2.0","error":{"code":-32602,"message":"Invalid subscription id."},"id":1}"#.into()),
-		);
+		let (response, _) = rpc.raw_json_request(&unsub_req).await.unwrap();
+		let expected = r#"{"jsonrpc":"2.0","result":false,"id":1}"#;
+
+		assert_eq!(response, expected);
 	}

-	#[test]
-	fn subscribe_and_unsubscribe_with_wrong_id() {
-		let (io, _) = setup_io_handler();
-		let (meta, _) = setup_session();
-
-		// Subscribe
-		let sub_request =
-			r#"{"jsonrpc":"2.0","method":"beefy_subscribeJustifications","params":[],"id":1}"#;
-		let resp = io.handle_request_sync(sub_request, meta.clone());
-		let resp: Output = serde_json::from_str(&resp.unwrap()).unwrap();
-		assert!(matches!(resp, Output::Success(_)));
+	#[tokio::test]
+	async fn subscribe_and_unsubscribe_with_wrong_id() {
+		let (rpc, _) = setup_io_handler();
+		// Subscribe call.
+		let _sub = rpc
+			.subscribe("beefy_subscribeJustifications", EmptyParams::new())
+			.await
+			.unwrap();

 		// Unsubscribe with wrong ID
-		assert_eq!(
-			io.handle_request_sync(
+		let (response, _) = rpc
+			.raw_json_request(
 				r#"{"jsonrpc":"2.0","method":"beefy_unsubscribeJustifications","params":["FOO"],"id":1}"#,
-				meta.clone()
-			),
-			Some(r#"{"jsonrpc":"2.0","error":{"code":-32602,"message":"Invalid subscription id."},"id":1}"#.into())
-		);
+			)
+			.await
+			.unwrap();
+		let expected = r#"{"jsonrpc":"2.0","result":false,"id":1}"#;
+
+		assert_eq!(response, expected);
 	}

 	fn create_commitment() -> BeefySignedCommitment<Block> {
@@ -357,18 +309,15 @@ mod tests {
 		}
 	}

-	#[test]
-	fn subscribe_and_listen_to_one_justification() {
-		let (io, commitment_sender) = setup_io_handler();
-		let (meta, receiver) = setup_session();
+	#[tokio::test]
+	async fn subscribe_and_listen_to_one_justification() {
+		let (rpc, commitment_sender) = setup_io_handler();

 		// Subscribe
-		let sub_request =
-			r#"{"jsonrpc":"2.0","method":"beefy_subscribeJustifications","params":[],"id":1}"#;
-
-		let resp = io.handle_request_sync(sub_request, meta.clone());
-		let mut resp: serde_json::Value = serde_json::from_str(&resp.unwrap()).unwrap();
-		let sub_id: String = serde_json::from_value(resp["result"].take()).unwrap();
+		let mut sub = rpc
+			.subscribe("beefy_subscribeJustifications", EmptyParams::new())
+			.await
+			.unwrap();

 		// Notify with commitment
 		let commitment = create_commitment();
@@ -376,21 +325,10 @@ mod tests {
 		r.unwrap();

 		// Inspect what we received
-		let recv = futures::executor::block_on(receiver.take(1).collect::<Vec<_>>());
-		let recv: Notification = serde_json::from_str(&recv[0]).unwrap();
-		let mut json_map = match recv.params {
-			Params::Map(json_map) => json_map,
-			_ => panic!(),
-		};
-
-		let recv_sub_id: String = serde_json::from_value(json_map["subscription"].take()).unwrap();
-		let recv_commitment: sp_core::Bytes =
-			serde_json::from_value(json_map["result"].take()).unwrap();
+		let (bytes, recv_sub_id) = sub.next::<sp_core::Bytes>().await.unwrap().unwrap();
 		let recv_commitment: BeefySignedCommitment<Block> =
-			Decode::decode(&mut &recv_commitment[..]).unwrap();
-
-		assert_eq!(recv.method, "beefy_justifications");
-		assert_eq!(recv_sub_id, sub_id);
+			Decode::decode(&mut &bytes[..]).unwrap();
+		assert_eq!(&recv_sub_id, sub.subscription_id());
 		assert_eq!(recv_commitment, commitment);
 	}
 }