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

Code before:
from test_framework.messages import COIN
from test_framework.util import (
    assert_equal,
    assert_raises_rpc_error,
    ensure_for,
)
from test_framework.wallet import MiniWallet
from test_framework.wallet_util import get_generate_key

START_HEIGHT = 199
SNAPSHOT_BASE_HEIGHT = 299
FINAL_HEIGHT = 399


class AssumeutxoTest(BitcoinTestFramework):
    def skip_test_if_missing_module(self):
        self.skip_if_no_wallet()

    def set_test_params(self):
        """Use the pregenerated, deterministic chain up to height 199."""
        self.num_nodes = 3
        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(3)
        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 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]

Review comment (suggestion):
nit, for smaller future diff if another node is added
```suggestion
            ["-fastprune", "-prune=1"],
```

Code after:
    ensure_for,
)
from test_framework.wallet import MiniWallet
from test_framework.wallet_util import get_generate_key

START_HEIGHT = 199
SNAPSHOT_BASE_HEIGHT = 299
FINAL_HEIGHT = 399


class AssumeutxoTest(BitcoinTestFramework):
    def skip_test_if_missing_module(self):
        self.skip_if_no_wallet()

    def set_test_params(self):
        """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)