From f3b3e6243377a919dd48305b90d72fbcef4f9f07 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bastian=20K=C3=B6cher?= <bkchr@users.noreply.github.com>
Date: Tue, 30 Mar 2021 08:10:54 +0200
Subject: [PATCH] Get rid of `test-helpers` feature in sc-consensus-babe
 (#8486)

This is not really required and having a special `test-helpers` feature
is a bad idea anyway.
---
 bin/node/cli/Cargo.toml          |  2 +-
 bin/node/cli/src/service.rs      | 39 +++++++++++++++++++-------------
 client/consensus/babe/Cargo.toml |  3 ---
 client/consensus/babe/src/lib.rs | 38 -------------------------------
 4 files changed, 24 insertions(+), 58 deletions(-)

diff --git a/bin/node/cli/Cargo.toml b/bin/node/cli/Cargo.toml
index c30378e8fc836..9449edfbf6e08 100644
--- a/bin/node/cli/Cargo.toml
+++ b/bin/node/cli/Cargo.toml
@@ -116,7 +116,7 @@ sp-trie = { version = "3.0.0", default-features = false, path = "../../../primit
 [dev-dependencies]
 sc-keystore = { version = "3.0.0", path = "../../../client/keystore" }
 sc-consensus = { version = "0.9.0", path = "../../../client/consensus/common" }
-sc-consensus-babe = { version = "0.9.0", features = ["test-helpers"], path = "../../../client/consensus/babe" }
+sc-consensus-babe = { version = "0.9.0", path = "../../../client/consensus/babe" }
 sc-consensus-epochs = { version = "0.9.0", path = "../../../client/consensus/epochs" }
 sc-service-test = { version = "2.0.0", path = "../../../client/service/test" }
 futures = "0.3.9"
diff --git a/bin/node/cli/src/service.rs b/bin/node/cli/src/service.rs
index 1351782315be7..4dc99c1736910 100644
--- a/bin/node/cli/src/service.rs
+++ b/bin/node/cli/src/service.rs
@@ -638,27 +638,34 @@ mod tests {
 					None,
 				);
 
-				let epoch_descriptor = babe_link.epoch_changes().lock().epoch_descriptor_for_child_of(
-					descendent_query(&*service.client()),
-					&parent_hash,
-					parent_number,
-					slot.into(),
-				).unwrap().unwrap();
-
 				let mut digest = Digest::<H256>::default();
 
 				// even though there's only one authority some slots might be empty,
 				// so we must keep trying the next slots until we can claim one.
-				let babe_pre_digest = loop {
-					inherent_data.replace_data(sp_timestamp::INHERENT_IDENTIFIER, &(slot * SLOT_DURATION));
-					if let Some(babe_pre_digest) = sc_consensus_babe::test_helpers::claim_slot(
+				let (babe_pre_digest, epoch_descriptor) = loop {
+					inherent_data.replace_data(
+						sp_timestamp::INHERENT_IDENTIFIER,
+						&(slot * SLOT_DURATION),
+					);
+
+					let epoch_descriptor = babe_link.epoch_changes().lock().epoch_descriptor_for_child_of(
+						descendent_query(&*service.client()),
+						&parent_hash,
+						parent_number,
+						slot.into(),
+					).unwrap().unwrap();
+
+					let epoch = babe_link.epoch_changes().lock().epoch_data(
+						&epoch_descriptor,
+						|slot| sc_consensus_babe::Epoch::genesis(&babe_link.config(), slot),
+					).unwrap();
+
+					if let Some(babe_pre_digest) = sc_consensus_babe::authorship::claim_slot(
 						slot.into(),
-						&parent_header,
-						&*service.client(),
-						keystore.clone(),
-						&babe_link,
-					) {
-						break babe_pre_digest;
+						&epoch,
+						&keystore,
+					).map(|(digest, _)| digest) {
+						break (babe_pre_digest, epoch_descriptor)
 					}
 
 					slot += 1;
diff --git a/client/consensus/babe/Cargo.toml b/client/consensus/babe/Cargo.toml
index 14d48fba1bb57..cd989650671d1 100644
--- a/client/consensus/babe/Cargo.toml
+++ b/client/consensus/babe/Cargo.toml
@@ -65,6 +65,3 @@ substrate-test-runtime-client = { version = "2.0.0", path = "../../../test-utils
 sc-block-builder = { version = "0.9.0", path = "../../block-builder" }
 rand_chacha = "0.2.2"
 tempfile = "3.1.0"
-
-[features]
-test-helpers = []
diff --git a/client/consensus/babe/src/lib.rs b/client/consensus/babe/src/lib.rs
index 28cfec1238e42..8fd41aa98014c 100644
--- a/client/consensus/babe/src/lib.rs
+++ b/client/consensus/babe/src/lib.rs
@@ -1647,41 +1647,3 @@ pub fn import_queue<Block: BlockT, Client, SelectChain, Inner, CAW>(
 		registry,
 	))
 }
-
-/// BABE test helpers. Utility methods for manually authoring blocks.
-#[cfg(feature = "test-helpers")]
-pub mod test_helpers {
-	use super::*;
-
-	/// Try to claim the given slot and return a `BabePreDigest` if
-	/// successful.
-	pub fn claim_slot<B, C>(
-		slot: Slot,
-		parent: &B::Header,
-		client: &C,
-		keystore: SyncCryptoStorePtr,
-		link: &BabeLink<B>,
-	) -> Option<PreDigest> where
-		B: BlockT,
-		C: ProvideRuntimeApi<B> +
-			ProvideCache<B> +
-			HeaderBackend<B> +
-			HeaderMetadata<B, Error = ClientError>,
-		C::Api: BabeApi<B>,
-	{
-		let epoch_changes = link.epoch_changes.lock();
-		let epoch = epoch_changes.epoch_data_for_child_of(
-			descendent_query(client),
-			&parent.hash(),
-			parent.number().clone(),
-			slot,
-			|slot| Epoch::genesis(&link.config, slot),
-		).unwrap().unwrap();
-
-		authorship::claim_slot(
-			slot,
-			&epoch,
-			&keystore,
-		).map(|(digest, _)| digest)
-	}
-}