From 3c7111d82ddb23aa2b288501ffd0081c5a691e6a Mon Sep 17 00:00:00 2001
From: Guillaume Potier <elmattic@users.noreply.github.com>
Date: Mon, 22 Apr 2024 12:16:56 +0200
Subject: [PATCH] Fix balance comparison to zero test (#4240)

---
 CHANGELOG.md                           |  4 +++
 scripts/tests/calibnet_wallet_check.sh | 15 ++++++----
 src/wallet/subcommands/wallet_cmd.rs   | 41 ++++++++++++++++++--------
 3 files changed, 42 insertions(+), 18 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c003a367cde2..ee020038e7a7 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -58,6 +58,10 @@
 - [#4183](https://github.com/ChainSafe/forest/issues/4183) Add support for the
   `Filecoin.EthGetBlockByNumber` RPC method.
 
+- [#4240](https://github.com/ChainSafe/forest/pull/4240) Added `--fixed-unit`
+  and `--exact-balance` flags to `forest-wallet balance` similarly to
+  `forest-wallet list` subcommand.
+
 ### Changed
 
 - [#4170](https://github.com/ChainSafe/forest/pull/4170) Change the default
diff --git a/scripts/tests/calibnet_wallet_check.sh b/scripts/tests/calibnet_wallet_check.sh
index 1c8c98ff3988..c17b4110737a 100755
--- a/scripts/tests/calibnet_wallet_check.sh
+++ b/scripts/tests/calibnet_wallet_check.sh
@@ -28,6 +28,9 @@ forest_init
 # Amount to send to 2nd address (note: `send` command defaults to FIL if no units are specified)
 FIL_AMT="500 atto FIL"
 
+# Amount for an empty wallet
+FIL_ZERO="0 FIL"
+
 $FOREST_WALLET_PATH import preloaded_wallet.key
 $FOREST_WALLET_PATH --remote-wallet import preloaded_wallet.key
 
@@ -77,24 +80,24 @@ MSG=$($FOREST_WALLET_PATH send "$ADDR_TWO" "$FIL_AMT")
 MSG_REMOTE=$($FOREST_WALLET_PATH --remote-wallet send "$ADDR_THREE" "$FIL_AMT")
 : "$MSG_REMOTE"
 
-ADDR_TWO_BALANCE=0
+ADDR_TWO_BALANCE=$FIL_ZERO
 i=0
-while [[ $i != 20 && $ADDR_TWO_BALANCE == 0 ]]; do
+while [[ $i != 20 && $ADDR_TWO_BALANCE == "$FIL_ZERO" ]]; do
   i=$((i+1))
   
   : "Checking balance $i/20"
   sleep 30s
-  ADDR_TWO_BALANCE=$($FOREST_WALLET_PATH balance "$ADDR_TWO")
+  ADDR_TWO_BALANCE=$($FOREST_WALLET_PATH balance "$ADDR_TWO" -e)
 done
 
-ADDR_THREE_BALANCE=0
+ADDR_THREE_BALANCE=$FIL_ZERO
 i=0
-while [[ $i != 20 && $ADDR_THREE_BALANCE == 0 ]]; do
+while [[ $i != 20 && $ADDR_THREE_BALANCE == "$FIL_ZERO" ]]; do
   i=$((i+1))
 
   : "Checking balance $i/20"
   sleep 30s
-  ADDR_THREE_BALANCE=$($FOREST_WALLET_PATH --remote-wallet balance "$ADDR_TWO")
+  ADDR_THREE_BALANCE=$($FOREST_WALLET_PATH --remote-wallet balance "$ADDR_THREE" -e)
 done
 
 # wallet list should contain address two with transfered FIL amount
diff --git a/src/wallet/subcommands/wallet_cmd.rs b/src/wallet/subcommands/wallet_cmd.rs
index 8e0f9318b327..de09996233b1 100644
--- a/src/wallet/subcommands/wallet_cmd.rs
+++ b/src/wallet/subcommands/wallet_cmd.rs
@@ -209,6 +209,15 @@ pub enum WalletCommands {
     Balance {
         /// The address of the account to check
         address: String,
+        /// Output is rounded to 4 significant figures by default.
+        /// Do not round
+        // ENHANCE(aatifsyed): add a --round/--no-round argument pair
+        #[arg(long, alias = "exact-balance", short_alias = 'e')]
+        no_round: bool,
+        /// Output may be given an SI prefix like `atto` by default.
+        /// Do not do this, showing whole FIL at all times.
+        #[arg(long, alias = "fixed-unit", short_alias = 'f')]
+        no_abbrev: bool,
     },
     /// Get the default address of the wallet
     Default,
@@ -313,11 +322,15 @@ impl WalletCommands {
                 println!("{addr}");
                 Ok(())
             }
-            Self::Balance { address } => {
+            Self::Balance {
+                address,
+                no_round,
+                no_abbrev,
+            } => {
                 let StrictAddress(address) = StrictAddress::from_str(&address)
                     .with_context(|| format!("Invalid address: {address}"))?;
                 let balance = WalletBalance::call(&backend.remote, (address,)).await?;
-                println!("{balance}");
+                println!("{}", format_balance(&balance, no_round, no_abbrev));
                 Ok(())
             }
             Self::Default => {
@@ -405,16 +418,7 @@ impl WalletCommands {
                     let balance_token_amount =
                         WalletBalance::call(&backend.remote, (address,)).await?;
 
-                    let balance_string = match (no_round, no_abbrev) {
-                        // no_round, absolute
-                        (true, true) => format!("{:#}", balance_token_amount.pretty()),
-                        // no_round, relative
-                        (true, false) => format!("{}", balance_token_amount.pretty()),
-                        // round, absolute
-                        (false, true) => format!("{:#.4}", balance_token_amount.pretty()),
-                        // round, relative
-                        (false, false) => format!("{:.4}", balance_token_amount.pretty()),
-                    };
+                    let balance_string = format_balance(&balance_token_amount, no_round, no_abbrev);
 
                     println!("{address:41}  {default_address_mark:7}  {balance_string}");
                 }
@@ -561,3 +565,16 @@ fn input_password_to_load_encrypted_keystore(data_dir: PathBuf) -> dialoguer::Re
         .into_inner()
         .expect("validation succeeded, so keystore must be emplaced"))
 }
+
+fn format_balance(balance: &TokenAmount, no_round: bool, no_abbrev: bool) -> String {
+    match (no_round, no_abbrev) {
+        // no_round, absolute
+        (true, true) => format!("{:#}", balance.pretty()),
+        // no_round, relative
+        (true, false) => format!("{}", balance.pretty()),
+        // round, absolute
+        (false, true) => format!("{:#.4}", balance.pretty()),
+        // round, relative
+        (false, false) => format!("{:.4}", balance.pretty()),
+    }
+}