group related wallet CLI commands together

mintlayer · Jan 5, 2024 · 2119ca3 · 2119ca3
1 parent 768351b
commit 2119ca3
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 42 deletions.
diff --git a/test/functional/test_framework/wallet_cli_controller.py b/test/functional/test_framework/wallet_cli_controller.py
@@ -272,10 +272,10 @@ async def sync(self) -> str:
         return await self._write_command("syncwallet\n")
 
     async def start_staking(self) -> str:
-        return await self._write_command(f"startstaking\n")
+        return await self._write_command(f"staking start\n")
 
     async def stop_staking(self) -> str:
-        return await self._write_command(f"stopstaking\n")
+        return await self._write_command(f"staking stop\n")
 
     async def get_addresses_usage(self) -> str:
         return await self._write_command("showreceiveaddresses\n")

diff --git a/wallet/wallet-cli-lib/src/commands/mod.rs b/wallet/wallet-cli-lib/src/commands/mod.rs
@@ -18,7 +18,7 @@ mod helper_types;
 use std::{path::PathBuf, str::FromStr, sync::Arc};
 
 use chainstate::TokenIssuanceError;
-use clap::Parser;
+use clap::{Parser, Subcommand};
 use common::{
     address::Address,
     chain::{
@@ -53,6 +53,18 @@ use self::helper_types::{
     CliWithLocked,
 };
 
+#[derive(Subcommand, Clone, Debug)]
+pub enum StakingCommand {
+    /// Start staking
+    Start,
+
+    /// Stop staking
+    Stop,
+
+    /// Pool balance
+    PoolBalance { pool_id: String },
+}
+
 #[derive(Debug, Parser)]
 #[clap(rename_all = "lower")]
 pub enum WalletCommand {
@@ -151,14 +163,8 @@ pub enum WalletCommand {
         account_index: U31,
     },
 
-    /// Start staking
-    StartStaking,
-
-    StopStaking,
-
-    StakePoolBalance {
-        pool_id: String,
-    },
+    #[clap(subcommand)]
+    Staking(StakingCommand),
 
     /// Submit a block to be included in the chain
     ///
@@ -763,35 +769,8 @@ impl CommandHandler {
                 })
             }
 
-            WalletCommand::StartStaking => {
-                self.get_synced_controller()
-                    .await?
-                    .start_staking()
-                    .map_err(WalletCliError::Controller)?;
-                Ok(ConsoleCommand::Print(
-                    "Staking started successfully".to_owned(),
-                ))
-            }
-
-            WalletCommand::StopStaking => {
-                let (controller, selected_account) = self.get_controller_and_selected_acc()?;
-                controller.stop_staking(selected_account).map_err(WalletCliError::Controller)?;
-                Ok(ConsoleCommand::Print("Success".to_owned()))
-            }
-
-            WalletCommand::StakePoolBalance { pool_id } => {
-                let pool_id = parse_pool_id(chain_config, pool_id.as_str())?;
-                let balance_opt = rpc_client
-                    .get_stake_pool_balance(pool_id)
-                    .await
-                    .map_err(WalletCliError::RpcError)?;
-                match balance_opt {
-                    Some(balance) => Ok(ConsoleCommand::Print(print_coin_amount(
-                        chain_config,
-                        balance,
-                    ))),
-                    None => Ok(ConsoleCommand::Print("Not found".to_owned())),
-                }
+            WalletCommand::Staking(command) => {
+                self.handle_staking_command(command, chain_config, rpc_client).await
             }
 
             WalletCommand::SubmitBlock { block } => {
@@ -1460,4 +1439,42 @@ impl CommandHandler {
             WalletCommand::ClearHistory => Ok(ConsoleCommand::ClearHistory),
         }
     }
+
+    async fn handle_staking_command(
+        &mut self,
+        command: StakingCommand,
+        chain_config: &Arc<ChainConfig>,
+        rpc_client: &NodeRpcClient,
+    ) -> Result<ConsoleCommand, WalletCliError> {
+        match command {
+            StakingCommand::Start => {
+                self.get_synced_controller()
+                    .await?
+                    .start_staking()
+                    .map_err(WalletCliError::Controller)?;
+                Ok(ConsoleCommand::Print(
+                    "Staking started successfully".to_owned(),
+                ))
+            }
+            StakingCommand::Stop => {
+                let (controller, selected_account) = self.get_controller_and_selected_acc()?;
+                controller.stop_staking(selected_account).map_err(WalletCliError::Controller)?;
+                Ok(ConsoleCommand::Print("Success".to_owned()))
+            }
+            StakingCommand::PoolBalance { pool_id } => {
+                let pool_id = parse_pool_id(chain_config, pool_id.as_str())?;
+                let balance_opt = rpc_client
+                    .get_stake_pool_balance(pool_id)
+                    .await
+                    .map_err(WalletCliError::RpcError)?;
+                match balance_opt {
+                    Some(balance) => Ok(ConsoleCommand::Print(print_coin_amount(
+                        chain_config,
+                        balance,
+                    ))),
+                    None => Ok(ConsoleCommand::Print("Not found".to_owned())),
+                }
+            }
+        }
+    }
 }
diff --git a/wallet/wallet-cli-lib/src/lib.rs b/wallet/wallet-cli-lib/src/lib.rs
@@ -25,7 +25,7 @@ type CliController = wallet_controller::RpcController<wallet::wallet_events::Wal
 use std::sync::Arc;
 
 use cli_event_loop::Event;
-use commands::WalletCommand;
+use commands::{StakingCommand, WalletCommand};
 use common::chain::{
     config::{regtest_options::regtest_chain_config, ChainType},
     ChainConfig,
@@ -137,7 +137,7 @@ pub async fn run(
         let (res_tx, res_rx) = tokio::sync::oneshot::channel();
         event_tx
             .send(Event::HandleCommand {
-                command: WalletCommand::StartStaking,
+                command: WalletCommand::Staking(StakingCommand::Start),
                 res_tx,
             })
             .expect("should not fail");