test: Check that a failed wallet migration is cleaned up

Author: achow101

test/functional/wallet_migration.py

@@ -6,11 +6,14 @@
 
 import random
 import shutil
+import struct
+
 from test_framework.address import (
     script_to_p2sh,
     key_to_p2pkh,
     key_to_p2wpkh,
 )
+from test_framework.bdb import BTREE_MAGIC
 from test_framework.descriptors import descsum_create
 from test_framework.key import ECPubKey
 from test_framework.test_framework import BitcoinTestFramework
@@ -20,6 +23,7 @@
     assert_equal,
     assert_raises_rpc_error,
     find_vout_for_address,
+    sha256sum_file,
 )
 from test_framework.wallet_util import (
     get_generate_key,
@@ -827,6 +831,43 @@ def test_hybrid_pubkey(self):
 
         wallet.unloadwallet()
 
+    def test_failed_migration_cleanup(self):
+        self.log.info("Test that a failed migration is cleaned up")
+        wallet = self.create_legacy_wallet("failed")
+
+        # Make a copy of the wallet with the solvables wallet name so that we are unable
+        # to create the solvables wallet when migrating, thus failing to migrate
+        wallet.unloadwallet()
+        solvables_path = self.nodes[0].wallets_path / "failed_solvables"
+        shutil.copytree(self.nodes[0].wallets_path / "failed", solvables_path)
+        original_shasum = sha256sum_file(solvables_path / "wallet.dat")
+
+        self.nodes[0].loadwallet("failed")
+
+        # Add a multisig so that a solvables wallet is created
+        wallet.addmultisigaddress(2, [wallet.getnewaddress(), get_generate_key().pubkey])
+        wallet.importaddress(get_generate_key().p2pkh_addr)
+
+        assert_raises_rpc_error(-4, "Failed to create new watchonly wallet", wallet.migratewallet)
+
+        assert "failed" in self.nodes[0].listwallets()
+        assert "failed_watchonly" not in self.nodes[0].listwallets()
+        assert "failed_solvables" not in self.nodes[0].listwallets()
+
+        assert not (self.nodes[0].wallets_path / "failed_watchonly").exists()
+        # Since the file in failed_solvables is one that we put there, migration shouldn't touch it
+        assert solvables_path.exists()
+        new_shasum = sha256sum_file(solvables_path / "wallet.dat")
+        assert_equal(original_shasum, new_shasum)
+
+        wallet.unloadwallet()
+        # Check the wallet we tried to migrate is still BDB
+        with open(self.nodes[0].wallets_path / "failed" / "wallet.dat", "rb") as f:
+            data = f.read(16)
+            _, _, magic = struct.unpack("QII", data)
+            assert_equal(magic, BTREE_MAGIC)
+
+
     def run_test(self):
         self.generate(self.nodes[0], 101)
 
@@ -845,6 +886,7 @@ def run_test(self):
         self.test_migrate_raw_p2sh()
         self.test_conflict_txs()
         self.test_hybrid_pubkey()
+        self.test_failed_migration_cleanup()
 
 if __name__ == '__main__':
     WalletMigrationTest().main()