Repository: bitcoin/bitcoin
File: test/functional/wallet_assumeutxo.py (Python)

Code before:
    def add_options(self, parser):
        self.add_wallet_options(parser, legacy=False)

    def set_test_params(self):
        """Use the pregenerated, deterministic chain up to height 199."""
        self.num_nodes = 2
        self.rpc_timeout = 120
        self.extra_args = [
            [],
            [],
        ]

    def setup_network(self):
        """Start with the nodes disconnected so that one can generate a snapshot
        including blocks the other hasn't yet seen."""
        self.add_nodes(2)
        self.start_nodes(extra_args=self.extra_args)

    def test_descriptor_import(self, node, wallet_name, key, timestamp, expected_error_message=None):
        import_request = [{"desc": descsum_create("pkh(" + key.pubkey + ")"),
                           "timestamp": timestamp,
                           "label": "Descriptor import test"}]
        wrpc = node.get_wallet_rpc(wallet_name)
        result = wrpc.importdescriptors(import_request)

        if expected_error_message is None:
            assert_equal(result[0]['success'], True)
        else:
            assert_equal(result[0]['error']['code'], -1)
            assert_equal(result[0]['error']['message'], expected_error_message)

    def run_test(self):
        """
        Bring up two (disconnected) nodes, mine some new blocks on the first,
        and generate a UTXO snapshot.

        Load the snapshot into the second, ensure it syncs to tip and completes
        background validation when connected to the first.
        """
        n0 = self.nodes[0]
        n1 = self.nodes[1]

        self.mini_wallet = MiniWallet(n0)

        # Mock time for a deterministic chain
        for n in self.nodes:
            n.setmocktime(n.getblockheader(n.getbestblockhash())['time'])

        n0.createwallet('w')
        w = n0.get_wallet_rpc("w")


Review comment (refactor):
You can simplify this by just returning the result from the function and then doing these checks respectively in the two places where the function is used now. Should drop the `test_` from the name as well then.

Code after:
        """Use the pregenerated, deterministic chain up to height 199."""
        self.num_nodes = 4
        self.rpc_timeout = 120
        self.extra_args = [
            [],
            [],
            [],
            ["-fastprune", "-prune=1"],
        ]

    def setup_network(self):
        """Start with the nodes disconnected so that one can generate a snapshot
        including blocks the other hasn't yet seen."""
        self.add_nodes(4)
        self.start_nodes(extra_args=self.extra_args)

    def import_descriptor(self, node, wallet_name, key, timestamp):
        import_request = [{"desc": descsum_create("pkh(" + key.pubkey + ")"),
                           "timestamp": timestamp,
                           "label": "Descriptor import test"}]
        wrpc = node.get_wallet_rpc(wallet_name)
        return wrpc.importdescriptors(import_request)

    def validate_snapshot_import(self, node, loaded, base_hash):
        assert_equal(loaded['coins_loaded'], SNAPSHOT_BASE_HEIGHT)
        assert_equal(loaded['base_height'], SNAPSHOT_BASE_HEIGHT)

        normal, snapshot = node.getchainstates()["chainstates"]
        assert_equal(normal['blocks'], START_HEIGHT)
        assert 'snapshot_blockhash' not in normal
        assert_equal(normal['validated'], True)
        assert_equal(snapshot['blocks'], SNAPSHOT_BASE_HEIGHT)
        assert_equal(snapshot['snapshot_blockhash'], base_hash)
        assert_equal(snapshot['validated'], False)

        assert_equal(node.getblockchaininfo()["blocks"], SNAPSHOT_BASE_HEIGHT)

    def complete_background_validation(self, node):
        self.connect_nodes(0, node.index)

        # Ensuring snapshot chain syncs to tip
        self.wait_until(lambda: node.getchainstates()['chainstates'][-1]['blocks'] == FINAL_HEIGHT)
        self.sync_blocks(nodes=(self.nodes[0], node))

        # Ensuring background validation completes
        self.wait_until(lambda: len(node.getchainstates()['chainstates']) == 1)

    def test_backup_during_background_sync_pruned_node(self, n3, dump_output, expected_error_message):
        self.log.info("Backup from the snapshot height can be loaded during background sync (pruned node)")
        loaded = n3.loadtxoutset(dump_output['path'])
        assert_greater_than(n3.pruneblockchain(START_HEIGHT), 0)