accidental removal reverted

utils/frame/try-runtime/cli/src/commands/follow_chain.rs

@@ -20,20 +20,16 @@ use crate::{
 	state_machine_call_with_proof, SharedParams, LOG_TARGET,
 };
 use jsonrpsee::{
-	core::{
-		async_trait,
-		client::{Client, Subscription, SubscriptionClientT},
-	},
+	core::client::{Subscription, SubscriptionClientT},
 	ws_client::WsClientBuilder,
 };
 use parity_scale_codec::{Decode, Encode};
 use remote_externalities::{rpc_api, Builder, Mode, OnlineConfig};
 use sc_executor::NativeExecutionDispatch;
 use sc_service::Configuration;
-use serde::de::DeserializeOwned;
 use sp_core::H256;
-use sp_runtime::traits::{Block as BlockT, Header as HeaderT, NumberFor};
-use std::{collections::VecDeque, fmt::Debug, marker::PhantomData, str::FromStr};
+use sp_runtime::traits::{Block as BlockT, Header, NumberFor};
+use std::{fmt::Debug, str::FromStr};
 
 const SUB: &str = "chain_subscribeFinalizedHeads";
 const UN_SUB: &str = "chain_unsubscribeFinalizedHeads";
@@ -235,6 +231,14 @@ where
 
 	let mut rpc_service = rpc_api::RpcService::new(&command.uri, command.keep_connection);
 
+	let header_provider: RpcHeaderProvider<Block> =
+		RpcHeaderProvider { uri: command.uri.clone(), _phantom: PhantomData {} };
+	let mut finalized_headers: FinalizedHeaders<
+		Block,
+		RpcHeaderProvider<Block>,
+		Subscription<Block::Header>,
+	> = FinalizedHeaders::new(header_provider, subscription);
+
 	while let Some(header) = finalized_headers.next().await {
 		let hash = header.hash();
 		let number = header.number();
@@ -326,74 +330,3 @@ where
 	log::error!(target: LOG_TARGET, "ws subscription must have terminated.");
 	Ok(())
 }
-
-#[cfg(test)]
-mod tests {
-	use super::*;
-	use sp_runtime::testing::{Block as TBlock, ExtrinsicWrapper, Header};
-
-	type Block = TBlock<ExtrinsicWrapper<()>>;
-	type BlockNumber = u64;
-	type Hash = H256;
-
-	struct MockHeaderProvider(pub VecDeque<BlockNumber>);
-
-	fn headers() -> Vec<Header> {
-		let mut headers = vec![Header::new_from_number(0)];
-		for n in 1..11 {
-			headers.push(Header {
-				parent_hash: headers.last().unwrap().hash(),
-				..Header::new_from_number(n)
-			})
-		}
-		headers
-	}
-
-	#[async_trait]
-	impl HeaderProvider<Block> for MockHeaderProvider {
-		async fn get_header(&mut self, _hash: Hash) -> Header {
-			let height = self.0.pop_front().unwrap();
-			headers()[height as usize].clone()
-		}
-	}
-
-	struct MockHeaderSubscription(pub VecDeque<BlockNumber>);
-
-	#[async_trait]
-	impl HeaderSubscription<Block> for MockHeaderSubscription {
-		async fn next_header(&mut self) -> Option<Header> {
-			self.0.pop_front().map(|h| headers()[h as usize].clone())
-		}
-	}
-
-	#[tokio::test]
-	async fn finalized_headers_works_when_every_block_comes_from_subscription() {
-		let heights = vec![4, 5, 6, 7];
-
-		let provider = MockHeaderProvider(vec![].into());
-		let subscription = MockHeaderSubscription(heights.clone().into());
-		let mut headers = FinalizedHeaders::new(provider, subscription);
-
-		for h in heights {
-			assert_eq!(h, headers.next().await.unwrap().number);
-		}
-		assert_eq!(None, headers.next().await);
-	}
-
-	#[tokio::test]
-	async fn finalized_headers_come_from_subscription_and_provider_if_in_need() {
-		let all_heights = 3..11;
-		let heights_in_subscription = vec![3, 4, 6, 10];
-		// Consecutive headers will be requested in the reversed order.
-		let heights_not_in_subscription = vec![5, 9, 8, 7];
-
-		let provider = MockHeaderProvider(heights_not_in_subscription.into());
-		let subscription = MockHeaderSubscription(heights_in_subscription.into());
-		let mut headers = FinalizedHeaders::new(provider, subscription);
-
-		for h in all_heights {
-			assert_eq!(h, headers.next().await.unwrap().number);
-		}
-		assert_eq!(None, headers.next().await);
-	}
-}