Add test_choose_option_question

safe_cli/operators/hw_accounts/ledger_manager.py

@@ -103,7 +103,7 @@ def sign_eip712(self, safe_tx: SafeTx, accounts: List[LedgerAccount]) -> SafeTx:
         for account in accounts:
             print_formatted_text(
                 HTML(
-                    "Ensure to compare in your ledger before to sign that domain_hash and message_hash are  both correct"
+                    "<ansired>Make sure in your ledger before signing that domain_hash and message_hash are both correct</ansired>"
                 )
             )
             print_formatted_text(HTML(f"Domain_hash: <b>{domain_hash.hex()}</b>"))

safe_cli/operators/safe_operator.py

@@ -292,16 +292,18 @@ def load_ledger_cli_owners(self, legacy_account: bool = False):
         option = choose_option_question(
             "Select the owner address", len(ledger_accounts) - 1
         )
+        if option is None:
+            return None
         address, derivation_path = ledger_accounts[option]
-        if self.ledger_manager.add_account(derivation_path):
-            balance = self.ethereum_client.get_balance(address)
-            print_formatted_text(
-                HTML(
-                    f"Loaded account <b>{address}</b> "
-                    f'with balance={Web3.from_wei(balance, "ether")} ether'
-                    f"Ledger account cannot be defined as sender"
-                )
+        self.ledger_manager.add_account(derivation_path)
+        balance = self.ethereum_client.get_balance(address)
+        print_formatted_text(
+            HTML(
+                f"Loaded account <b>{address}</b> "
+                f'with balance={Web3.from_wei(balance, "ether")} ether'
+                f"Ledger account cannot be defined as sender"
             )
+        )
 
     def unload_cli_owners(self, owners: List[str]):
         accounts_to_remove: Set[Account] = set()

safe_cli/utils.py

@@ -1,5 +1,7 @@
-import argparse
 import os
+from typing import Optional
+
+from prompt_toolkit import HTML, print_formatted_text
 
 from gnosis.eth import EthereumClient
 
@@ -46,12 +48,21 @@ def yes_or_no_question(question: str, default_no: bool = True) -> bool:
 
 def choose_option_question(
     question: str, number_options: int, default_option: int = 0
-) -> bool:
+) -> Optional[int]:
     if "PYTEST_CURRENT_TEST" in os.environ:
-        return 0  # Ignore confirmations when running tests
+        return default_option  # Ignore confirmations when running tests
     choices = f" [0-{number_options}] default {default_option}:"
-    reply = str(input(question + choices)).lower().strip() or str(default_option)
-    option = int(reply)
+    reply = str(get_input(question + choices)).lower().strip() or str(default_option)
+    try:
+        option = int(reply)
+    except ValueError:
+        print_formatted_text(HTML("<ansired> Option must be an integer </ansired>"))
+        return None
+
     if option not in range(0, number_options):
-        argparse.ArgumentTypeError(f"{option} is not between [0-{number_options}}}")
+        print_formatted_text(
+            HTML(f"<ansired> {option} is not between [0-{number_options}}} </ansired>")
+        )
+        return None
+
     return option

tests/test_utils.py

@@ -3,7 +3,7 @@
 from unittest import mock
 from unittest.mock import MagicMock
 
-from safe_cli.utils import yes_or_no_question
+from safe_cli.utils import choose_option_question, yes_or_no_question
 
 
 class TestUtils(unittest.TestCase):
@@ -34,6 +34,23 @@ def test_yes_or_no_question(self, input_mock: MagicMock):
 
         os.environ["PYTEST_CURRENT_TEST"] = pytest_current_test
 
+    @mock.patch("safe_cli.utils.get_input")
+    def test_choose_option_question(self, input_mock: MagicMock):
+        pytest_current_test = os.environ.pop("PYTEST_CURRENT_TEST")
+
+        input_mock.return_value = ""
+        self.assertEqual(choose_option_question("", 1), 0)
+        input_mock.return_value = ""
+        self.assertEqual(choose_option_question("", 5, 4), 4)
+        input_mock.return_value = "m"
+        self.assertIsNone(choose_option_question("", 1))
+        input_mock.return_value = "10"
+        self.assertIsNone(choose_option_question("", 1))
+        input_mock.return_value = "1"
+        self.assertEqual(choose_option_question("", 2), 1)
+
+        os.environ["PYTEST_CURRENT_TEST"] = pytest_current_test
+
 
 if __name__ == "__main__":
     unittest.main()