dragonx/qa/rpc-tests/wallet_persistence.py

#!/usr/bin/env python3
# Copyright (c) 2016-2024 The Hush developers
# Copyright (c) 2018 The Zcash developers
# Distributed under the GPLv3 software license, see the accompanying
# file COPYING or https://www.gnu.org/licenses/gpl-3.0.en.html

from test_framework.test_framework import BitcoinTestFramework
from test_framework.util import (
    assert_equal, assert_true,
    start_nodes, stop_nodes,
    initialize_chain_clean, connect_nodes_bi, wait_bitcoinds,
    wait_and_assert_operationid_status
)
from decimal import Decimal

def get_value_pool(pools, pool_id):
    """
    Return the valuePools entry with the given id, or None if this chain does
    not have that pool. DragonX's getblockchaininfo only reports the Sapling
    pool (no Sprout history exists on this chain), so the pools list can no
    longer be indexed positionally the way the upstream test did.
    """
    for pool in pools:
        if pool['id'] == pool_id:
            return pool
    return None

def assert_pool_values(pools, sprout_value, sapling_value):
    sprout = get_value_pool(pools, 'sprout')
    if sprout is not None:
        assert_equal(sprout['chainValue'], sprout_value)
    else:
        # No Sprout pool at all is the same statement as "the Sprout pool holds nothing"
        assert_equal(sprout_value, Decimal('0'))
    sapling = get_value_pool(pools, 'sapling')
    assert_true(sapling is not None, "Sapling value pool missing from getblockchaininfo")
    assert_equal(sapling['chainValue'], sapling_value)

class WalletPersistenceTest (BitcoinTestFramework):

    def setup_chain(self):
        print("Initializing test directory " + self.options.tmpdir)
        initialize_chain_clean(self.options.tmpdir, 3)

    def setup_network(self, split=False):
        # -listen=1/-bind: the framework's start_node() passes -connect=0, and DragonX (like
        # Bitcoin) reacts to -connect by soft-setting -listen=0. A non-listening node can never
        # accept the "addnode 127.0.0.1:<port>" that connect_nodes_bi() issues, so without this
        # the three nodes stay isolated and sync_all() spins forever. -bind keeps the listener on
        # loopback so a regtest node never becomes reachable from the public internet.
        self.nodes = start_nodes(3, self.options.tmpdir,
            extra_args=[[
                '-listen=1',
                '-bind=127.0.0.1',
                # -dns=0: DragonX appends node1..node10.dragonx.is to -addnode for every chain
                # named DRAGONX (hush_utils.h), and -connect=0 does not suppress -addnode. Without
                # this a regtest node dials the LIVE DragonX network and is fed mainnet headers.
                # The addnode calls connect_nodes_bi() makes use literal IPs, so they still work.
                '-dns=0',
                '-nuparams=5ba81b19:100', # Overwinter
                '-nuparams=76b809bb:201', # Sapling
            ]] * 3)
        connect_nodes_bi(self.nodes,0,1)
        connect_nodes_bi(self.nodes,1,2)
        self.is_network_split=False
        self.sync_all()

    def run_test(self):
        # Sanity-check the test harness
        self.nodes[0].generate(200)
        assert_equal(self.nodes[0].getblockcount(), 200)
        self.sync_all()

        # Verify Sapling address is persisted in wallet (even when Sapling is not yet active)
        sapling_addr = self.nodes[0].z_getnewaddress('sapling')

        # Make sure the node has the addresss
        addresses = self.nodes[0].z_listaddresses()
        assert_true(sapling_addr in addresses, "Should contain address before restart")

        # Restart the nodes
        stop_nodes(self.nodes)
        wait_bitcoinds()
        self.setup_network()

        # Make sure we still have the address after restarting
        addresses = self.nodes[0].z_listaddresses()
        assert_true(sapling_addr in addresses, "Should contain address after restart")

        # Activate Sapling
        self.nodes[0].generate(1)
        self.sync_all()

        # Node 0 shields funds to Sapling address
        taddr0 = self.nodes[0].getnewaddress()
        recipients = []
        recipients.append({"address": sapling_addr, "amount": Decimal('20')})
        myopid = self.nodes[0].z_sendmany(taddr0, recipients, 1, 0)
        wait_and_assert_operationid_status(self.nodes[0], myopid)

        self.sync_all()
        self.nodes[0].generate(1)
        self.sync_all()

        # Verify shielded balance
        assert_equal(self.nodes[0].z_getbalance(sapling_addr), Decimal('20'))

        # Verify size of shielded pools
        pools = self.nodes[0].getblockchaininfo()['valuePools']
        assert_pool_values(pools, Decimal('0'), Decimal('20'))

        # Restart the nodes
        stop_nodes(self.nodes)
        wait_bitcoinds()
        self.setup_network()

        # Verify size of shielded pools
        pools = self.nodes[0].getblockchaininfo()['valuePools']
        assert_pool_values(pools, Decimal('0'), Decimal('20'))

        # Node 0 sends some shielded funds to Node 1
        dest_addr = self.nodes[1].z_getnewaddress('sapling')
        recipients = []
        recipients.append({"address": dest_addr, "amount": Decimal('15')})
        myopid = self.nodes[0].z_sendmany(sapling_addr, recipients, 1, 0)
        wait_and_assert_operationid_status(self.nodes[0], myopid)

        self.sync_all()
        self.nodes[0].generate(1)
        self.sync_all()

        # Verify balances
        assert_equal(self.nodes[0].z_getbalance(sapling_addr), Decimal('5'))
        assert_equal(self.nodes[1].z_getbalance(dest_addr), Decimal('15'))

        # Restart the nodes
        stop_nodes(self.nodes)
        wait_bitcoinds()
        self.setup_network()

        # Verify balances
        assert_equal(self.nodes[0].z_getbalance(sapling_addr), Decimal('5'))
        assert_equal(self.nodes[1].z_getbalance(dest_addr), Decimal('15'))

        # Verify importing a spending key will update and persist the nullifiers and witnesses correctly
        sk0 = self.nodes[0].z_exportkey(sapling_addr)
        self.nodes[2].z_importkey(sk0, "yes")
        assert_equal(self.nodes[2].z_getbalance(sapling_addr), Decimal('5'))

        # Restart the nodes
        stop_nodes(self.nodes)
        wait_bitcoinds()
        self.setup_network()

        # Verify nullifiers persisted correctly by checking balance
        # Prior to PR #3590, there will be an error as spent notes are considered unspent:
        #    Assertion failed: expected: <25.00000000> but was: <5>
        assert_equal(self.nodes[2].z_getbalance(sapling_addr), Decimal('5'))

        # Verity witnesses persisted correctly by sending shielded funds
        recipients = []
        recipients.append({"address": dest_addr, "amount": Decimal('1')})
        myopid = self.nodes[2].z_sendmany(sapling_addr, recipients, 1, 0)
        wait_and_assert_operationid_status(self.nodes[2], myopid)

        self.sync_all()
        self.nodes[0].generate(1)
        self.sync_all()

        # Verify balances
        assert_equal(self.nodes[2].z_getbalance(sapling_addr), Decimal('4'))
        assert_equal(self.nodes[1].z_getbalance(dest_addr), Decimal('16'))

if __name__ == '__main__':
    WalletPersistenceTest().main()