refactor 'init_wallet_wizard' hook a bit. add test_wizard trezor test

SomberNight · SomberNight · commit b0464cc934ab · 2025-07-21T02:40:17.000Z
diff --git a/electrum/hw_wallet/plugin.py b/electrum/hw_wallet/plugin.py
@@ -39,6 +39,7 @@
     import threading
     from electrum.plugin import DeviceInfo
     from electrum.wallet import Abstract_Wallet
+    from electrum.wizard import AbstractWizard
 
 
 class HW_PluginBase(BasePlugin, ABC):
@@ -188,6 +189,14 @@ def wizard_entry_for_device(self, device_info: 'DeviceInfo', *, new_wallet: bool
         """
         pass
 
+    @hook
+    def init_wallet_wizard(self, wizard: 'AbstractWizard') -> None:
+        self.extend_wizard(wizard)
+
+    @abstractmethod
+    def extend_wizard(self, wizard: 'AbstractWizard') -> None:
+        pass
+
 
 class HardwareClientBase(ABC):
     handler = None  # type: Optional['HardwareHandlerBase']
diff --git a/electrum/plugins/bitbox02/qt.py b/electrum/plugins/bitbox02/qt.py
@@ -54,10 +54,6 @@ def on_button_click():
         device_name = "{} ({})".format(self.device, keystore.label)
         mpk_text.addButton(read_QIcon("eye1.png"), on_button_click, _("Show on {}").format(device_name))
 
-    @hook
-    def init_wallet_wizard(self, wizard: 'QENewWalletWizard'):
-        self.extend_wizard(wizard)
-
     # insert bitbox02 pages in new wallet wizard
     def extend_wizard(self, wizard: 'QENewWalletWizard'):
         super().extend_wizard(wizard)
diff --git a/electrum/plugins/coldcard/qt.py b/electrum/plugins/coldcard/qt.py
@@ -124,10 +124,6 @@ def show_settings_dialog(self, window, keystore):
         # - doesn't matter if device not connected, continue
         CKCCSettingsDialog(window, self, keystore).exec()
 
-    @hook
-    def init_wallet_wizard(self, wizard: 'QENewWalletWizard'):
-        self.extend_wizard(wizard)
-
     # insert coldcard pages in new wallet wizard
     def extend_wizard(self, wizard: 'QENewWalletWizard'):
         super().extend_wizard(wizard)
diff --git a/electrum/plugins/digitalbitbox/qt.py b/electrum/plugins/digitalbitbox/qt.py
@@ -47,10 +47,6 @@ def transaction_dialog_address_menu(self, menu, addr, wallet):
             return
         self._add_menu_action(menu, addr, wallet)
 
-    @hook
-    def init_wallet_wizard(self, wizard: 'QENewWalletWizard'):
-        self.extend_wizard(wizard)
-
     # insert digitalbitbox pages in new wallet wizard
     def extend_wizard(self, wizard: 'QENewWalletWizard'):
         super().extend_wizard(wizard)
diff --git a/electrum/plugins/jade/qt.py b/electrum/plugins/jade/qt.py
@@ -41,10 +41,6 @@ def transaction_dialog_address_menu(self, menu, addr, wallet):
             return
         self._add_menu_action(menu, addr, wallet)
 
-    @hook
-    def init_wallet_wizard(self, wizard: 'QENewWalletWizard'):
-        self.extend_wizard(wizard)
-
     # insert jade pages in new wallet wizard
     def extend_wizard(self, wizard: 'QENewWalletWizard'):
         super().extend_wizard(wizard)
diff --git a/electrum/plugins/keepkey/qt.py b/electrum/plugins/keepkey/qt.py
@@ -323,10 +323,6 @@ def create_handler(self, window):
     def pin_matrix_widget_class(self):
         return PinMatrixWidget
 
-    @hook
-    def init_wallet_wizard(self, wizard: 'QENewWalletWizard'):
-        self.extend_wizard(wizard)
-
     # insert keepkey pages in new wallet wizard
     def extend_wizard(self, wizard: 'QENewWalletWizard'):
         super().extend_wizard(wizard)
diff --git a/electrum/plugins/ledger/qt.py b/electrum/plugins/ledger/qt.py
@@ -40,10 +40,6 @@ def transaction_dialog_address_menu(self, menu, addr, wallet):
             return
         self._add_menu_action(menu, addr, wallet)
 
-    @hook
-    def init_wallet_wizard(self, wizard: 'QENewWalletWizard'):
-        self.extend_wizard(wizard)
-
     # insert ledger pages in new wallet wizard
     def extend_wizard(self, wizard: 'QENewWalletWizard'):
         super().extend_wizard(wizard)
diff --git a/electrum/plugins/safe_t/qt.py b/electrum/plugins/safe_t/qt.py
@@ -199,10 +199,6 @@ def create_handler(self, window):
     def pin_matrix_widget_class(self):
         return PinMatrixWidget
 
-    @hook
-    def init_wallet_wizard(self, wizard: 'QENewWalletWizard'):
-        self.extend_wizard(wizard)
-
     # insert safe_t pages in new wallet wizard
     def extend_wizard(self, wizard: 'QENewWalletWizard'):
         super().extend_wizard(wizard)
diff --git a/electrum/plugins/trezor/qt.py b/electrum/plugins/trezor/qt.py
@@ -466,10 +466,6 @@ def create_handler(self, window):
     def pin_matrix_widget_class(self):
         return PinMatrixWidget
 
-    @hook
-    def init_wallet_wizard(self, wizard: 'QENewWalletWizard'):
-        self.extend_wizard(wizard)
-
     # insert trezor pages in new wallet wizard
     def extend_wizard(self, wizard: 'QENewWalletWizard'):
         super().extend_wizard(wizard)
diff --git a/electrum/wizard.py b/electrum/wizard.py
@@ -172,7 +172,7 @@ def sanitize_stack_item(self, _stack_item) -> dict:
             "wallet_type", "keystore_type", "seed_variant", "seed_type", "seed_extend",
             "script_type", "derivation_path", "encrypt",
             # hardware devices:
-            "hardware_device", "hw_type", "label", "soft_device_id",
+            "hardware_device", "hw_type", "label", "soft_device_id", "xpub_encrypt",
             # inside keystore:
             "type", "pw_hash_version", "derivation", "root_fingerprint",
             # multisig:
@@ -285,8 +285,8 @@ def wallet_password_view(self, wizard_data: dict) -> str:
     def on_hardware_device(self, wizard_data: dict, new_wallet=True) -> str:
         current_cosigner = self.current_cosigner(wizard_data)
         _type, _info = current_cosigner['hardware_device']
-        run_hook('init_wallet_wizard', self)  # TODO: currently only used for hww, hook name might be confusing
         plugin = self.plugins.get_plugin(_type)
+        run_hook('init_wallet_wizard', self)  # TODO: currently only used for hww, hook name might be confusing
         return plugin.wizard_entry_for_device(_info, new_wallet=new_wallet)
 
     def validate_seed(self, seed: str, seed_variant: str, wallet_type: str) -> Tuple[bool, str, str, bool]:
diff --git a/tests/test_wizard.py b/tests/test_wizard.py
@@ -3,7 +3,7 @@
 from electrum import SimpleConfig
 from electrum.interface import ServerAddr
 from electrum.network import NetworkParameters, ProxySettings
-from electrum.plugin import Plugins
+from electrum.plugin import Plugins, DeviceInfo, Device
 from electrum.wizard import ServerConnectWizard, NewWalletWizard, WizardViewState
 from electrum.daemon import Daemon
 from electrum.wallet import Abstract_Wallet
@@ -49,6 +49,7 @@ def setUp(self):
         self.wallet_path = os.path.join(self.electrum_path, "somewallet")
         self.plugins = Plugins(self.config, gui_name='cmdline')
         self.plugins.load_plugin_by_name('trustedcoin')
+        # note: hw plugins are loaded on-demand
 
     def tearDown(self):
         self.plugins.stop()
@@ -479,3 +480,46 @@ async def test_2fa_haveseed_passphrase(self):
         v = w.resolve_next(v.view, d)
         self._set_password_and_check_address(v=v, w=w, recv_addr="bc1qcnu9ay4v3w0tawuxe6wlh6mh33rrpauqnufdgkxx7we8vpx3e6wqa25qud")
 
+    async def test_create_standard_wallet_trezor(self):
+        # bip39 seed for trezor: "history six okay anchor sheriff flock atom tomorrow foster aerobic eternal foam"
+        w = self._wizard_for(wallet_type='standard')
+        v = w._current
+        d = v.wizard_data
+        self.assertEqual('keystore_type', v.view)
+
+        d.update({'keystore_type': 'hardware'})
+        v = w.resolve_next(v.view, d)
+        self.assertEqual('choose_hardware_device', v.view)
+
+        d.update({
+            'hardware_device': (
+                'trezor',
+                DeviceInfo(
+                    device=Device(path='webusb:002:1', interface_number=-1, id_='webusb:002:1', product_key='Trezor', usage_page=0, transport_ui_string='webusb:002:1'),
+                    label='trezor_unittests', initialized=True, exception=None, plugin_name='trezor', soft_device_id='088C3F260B66F60E15DE0FA5', model_name='Trezor T'))})
+        v = w.resolve_next(v.view, d)
+        self.assertEqual('trezor_start', v.view)
+
+        d.update({'script_type': 'p2wpkh', 'derivation_path': 'm/84h/0h/0h'})
+        v = w.resolve_next(v.view, d)
+        self.assertEqual('trezor_xpub', v.view)
+
+        d.update({
+            'hw_type': 'trezor', 'master_key': 'zpub6qqp9XwsVMsovwzayXhFDJTpoc8VFoNy6mjkJHygou9NPRPDNR7MXVp9DM7qpacWwoePFWg7Gt5L5xnKNLmZYH8AFoTm2AAZA7LasycHu3n',
+            'root_fingerprint': '6306ee35', 'label': 'trezor_unittests', 'soft_device_id': '088C3F260B66F60E15DE0FA5',
+            'multisig_master_pubkey': 'zpub6qqp9XwsVMsovwzayXhFDJTpoc8VFoNy6mjkJHygou9NPRPDNR7MXVp9DM7qpacWwoePFWg7Gt5L5xnKNLmZYH8AFoTm2AAZA7LasycHu3n'})
+        v = w.resolve_next(v.view, d)
+        self.assertEqual('wallet_password_hardware', v.view)
+
+        d.update({'password': '03a580deb85ef85654ed177fc049867ce915a8b392a34a524123870925e48a5b9e', 'encrypt': True, 'xpub_encrypt': True})
+        self.assertTrue(w.is_last_view(v.view, d))
+        v = w.resolve_next(v.view, d)
+
+        wallet_path = os.path.join(w._daemon.config.get_datadir_wallet_path(), d['wallet_name'])
+        w.create_storage(wallet_path, d)
+
+        self.assertTrue(os.path.exists(wallet_path))
+        wallet = Daemon._load_wallet(wallet_path, password=d['password'], config=self.config)
+        self.assertEqual("bc1q7ltf4aq95rj695fu5aaa5mx5m9p55xyr2fy6y0", wallet.get_receiving_addresses()[0])
+        self.assertTrue(wallet.has_password())
+        self.assertTrue(wallet.has_storage_encryption())
