From e29538c1a79d1711b43addc9400d871f6aa32844 Mon Sep 17 00:00:00 2001
From: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>
Date: Tue, 3 Jan 2023 12:50:34 +0100
Subject: [PATCH] Add Clap wrapper for enum UpgradeCheckSelect

Signed-off-by: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>
---
 frame/support/src/traits/try_runtime.rs       | 15 -------
 .../cli/src/commands/on_runtime_upgrade.rs    | 39 ++++++++++++++-----
 2 files changed, 30 insertions(+), 24 deletions(-)

diff --git a/frame/support/src/traits/try_runtime.rs b/frame/support/src/traits/try_runtime.rs
index 273af09a145e5..29bce49762be5 100644
--- a/frame/support/src/traits/try_runtime.rs
+++ b/frame/support/src/traits/try_runtime.rs
@@ -106,21 +106,6 @@ impl UpgradeCheckSelect {
 	}
 }
 
-#[cfg(feature = "std")]
-impl core::str::FromStr for UpgradeCheckSelect {
-	type Err = &'static str;
-
-	fn from_str(s: &str) -> Result<Self, Self::Err> {
-		match s.to_lowercase().as_str() {
-			"none" => Ok(Self::None),
-			"all" => Ok(Self::All),
-			"pre-and-post" => Ok(Self::PreAndPost),
-			"try-state" => Ok(Self::TryState),
-			_ => Err("Invalid CheckSelector"),
-		}
-	}
-}
-
 /// Execute some checks to ensure the internal state of a pallet is consistent.
 ///
 /// Usually, these checks should check all of the invariants that are expected to be held on all of
diff --git a/utils/frame/try-runtime/cli/src/commands/on_runtime_upgrade.rs b/utils/frame/try-runtime/cli/src/commands/on_runtime_upgrade.rs
index 6f755244748df..38e4b8c23f6b7 100644
--- a/utils/frame/try-runtime/cli/src/commands/on_runtime_upgrade.rs
+++ b/utils/frame/try-runtime/cli/src/commands/on_runtime_upgrade.rs
@@ -16,7 +16,6 @@
 // limitations under the License.
 
 use crate::{build_executor, state_machine_call_with_proof, SharedParams, State, LOG_TARGET};
-use frame_try_runtime::UpgradeCheckSelect;
 use parity_scale_codec::{Decode, Encode};
 use sc_executor::sp_wasm_interface::HostFunctions;
 use sp_runtime::traits::{Block as BlockT, NumberFor};
@@ -30,18 +29,39 @@ pub struct OnRuntimeUpgradeCmd {
 	#[command(subcommand)]
 	pub state: State,
 
-	/// Select which optional checks to perform:
-	///
-	/// - `all`: Perform all checks (default).
-	/// - `pre-and-post`: Perform pre- and post-upgrade checks.
-	/// - `try-state`: Perform the try-state checks.
-	/// - `none`: Perform no checks.
+	/// Select which optional checks to perform.
 	///
 	/// Performing any checks will potentially invalidate the measured PoV/Weight.
-	#[clap(long, default_value = "All", verbatim_doc_comment)]
+	#[clap(long, value_enum, default_value_t = UpgradeCheckSelect::All)]
 	pub checks: UpgradeCheckSelect,
 }
 
+/// This is an adapter for [`frame_try_runtime::UpgradeCheckSelect`] since that does not implement
+/// `clap::ValueEnum`.
+#[derive(clap::ValueEnum, Debug, Clone, Copy)]
+#[value(rename_all = "kebab-case")]
+pub enum UpgradeCheckSelect {
+	/// Perform no checks.
+	None,
+	/// Perform all checks.
+	All,
+	/// Perform pre- and post-upgrade checks.
+	PreAndPost,
+	/// Perform the try-state checks.
+	TryState,
+}
+
+impl From<UpgradeCheckSelect> for frame_try_runtime::UpgradeCheckSelect {
+	fn from(x: UpgradeCheckSelect) -> Self {
+		match x {
+			UpgradeCheckSelect::None => Self::None,
+			UpgradeCheckSelect::All => Self::All,
+			UpgradeCheckSelect::PreAndPost => Self::PreAndPost,
+			UpgradeCheckSelect::TryState => Self::TryState,
+		}
+	}
+}
+
 pub(crate) async fn on_runtime_upgrade<Block, HostFns>(
 	shared: SharedParams,
 	command: OnRuntimeUpgradeCmd,
@@ -57,12 +77,13 @@ where
 {
 	let executor = build_executor(&shared);
 	let ext = command.state.into_ext::<Block, HostFns>(&shared, &executor, None).await?;
+	let checks: frame_try_runtime::UpgradeCheckSelect = command.checks.into();
 
 	let (_, encoded_result) = state_machine_call_with_proof::<Block, HostFns>(
 		&ext,
 		&executor,
 		"TryRuntime_on_runtime_upgrade",
-		command.checks.encode().as_ref(),
+		checks.encode().as_ref(),
 		Default::default(), // we don't really need any extensions here.
 		shared.export_proof,
 	)?;