diff --git a/qa/rpc-tests/wallet_accounts.py b/qa/rpc-tests/wallet_accounts.py index e8854ffef..79d92a518 100755 --- a/qa/rpc-tests/wallet_accounts.py +++ b/qa/rpc-tests/wallet_accounts.py @@ -40,6 +40,22 @@ class WalletAccountsTest(BitcoinTestFramework): assert_equal(expected[pool] * COIN, actual['pools'][pool]['valueZat']) assert_equal(actual['minimum_confirmations'], 1 if minconf is None else minconf) + # Check we only have balances in the expected pools. + # Remember that empty pools are omitted from the output. + def check_address_balance(self, address, expected, minconf=None): + if minconf is None: + actual = self.nodes[0].z_getbalanceforaddress(address) + else: + actual = self.nodes[0].z_getbalanceforaddress(address, minconf) + assert_equal(set(expected), set(actual['pools'])) + for pool in expected: + assert_equal(expected[pool] * COIN, actual['pools'][pool]['valueZat']) + assert_equal(actual['minimum_confirmations'], 1 if minconf is None else minconf) + + def check_balance(self, account, address, expected, minconf=None): + self.check_account_balance(account, expected, minconf) + self.check_address_balance(address, expected, minconf) + def run_test(self): # With a new wallet, the first account will be 0. account0 = self.nodes[0].z_getnewaccount() @@ -79,8 +95,8 @@ class WalletAccountsTest(BitcoinTestFramework): self.check_receiver_types(ua1, ['transparent', 'sapling']) # The balances of the accounts are all zero. - self.check_account_balance(0, {}) - self.check_account_balance(1, {}) + self.check_balance(0, ua0, {}) + self.check_balance(1, ua1, {}) # Manually send funds to one of the receivers in the UA. # TODO: Once z_sendmany supports UAs, receive to the UA instead of the receiver. @@ -98,14 +114,14 @@ class WalletAccountsTest(BitcoinTestFramework): # The new balance should not be visible with the default minconf, but should be # visible with minconf=0. self.sync_all() - self.check_account_balance(0, {}) - self.check_account_balance(0, {'sapling': 10}, 0) + self.check_balance(0, ua0, {}) + self.check_balance(0, ua0, {'sapling': 10}, 0) self.nodes[2].generate(1) self.sync_all() # The default minconf should now detect the balance. - self.check_account_balance(0, {'sapling': 10}) + self.check_balance(0, ua0, {'sapling': 10}) # Manually send funds from the UA receiver. # TODO: Once z_sendmany supports UAs, send from the UA instead of the receiver. @@ -125,8 +141,8 @@ class WalletAccountsTest(BitcoinTestFramework): # shown, as that transaction has been created and broadcast, and _might_ get mined # up until the transaction expires), or 9 (if we include the unmined transaction). self.sync_all() - self.check_account_balance(0, {}) - self.check_account_balance(0, {'sapling': 9}, 0) + self.check_balance(0, ua0, {}) + self.check_balance(0, ua0, {'sapling': 9}, 0) if __name__ == '__main__': diff --git a/src/wallet/rpcwallet.cpp b/src/wallet/rpcwallet.cpp index ad46851c3..6c29bad8c 100644 --- a/src/wallet/rpcwallet.cpp +++ b/src/wallet/rpcwallet.cpp @@ -3641,6 +3641,14 @@ UniValue z_getbalanceforaddress(const UniValue& params, bool fHelp) if (!fExperimentalOrchardWallet) { throw JSONRPCError(RPC_WALLET_ENCRYPTION_FAILED, "Error: the Orchard wallet experimental extensions are disabled."); } + + KeyIO keyIO(Params()); + auto decoded = keyIO.DecodePaymentAddress(params[0].get_str()); + if (!decoded.has_value()) { + throw JSONRPCError(RPC_INVALID_ADDRESS_OR_KEY, "Invalid address"); + } + auto address = decoded.value(); + int minconf = 1; if (params.size() > 1) { minconf = params[1].get_int(); @@ -3648,11 +3656,45 @@ UniValue z_getbalanceforaddress(const UniValue& params, bool fHelp) throw JSONRPCError(RPC_INVALID_PARAMETER, "Minimum number of confirmations cannot be less than 0"); } } + + LOCK(pwalletMain->cs_wallet); + + // Get the receivers for this address. + auto selector = pwalletMain->ToZTXOSelector(address, false); + if (!selector.has_value()) { + // The only way we'd reach this is if the address is a unified address for which + // we do not know its UFVK. + throw JSONRPCError( + RPC_INVALID_PARAMETER, + "Error: wallet does not have the Unified Full Viewing Key for the given address"); + } + + auto spendableInputs = pwalletMain->FindSpendableInputs(selector.value(), true, minconf); + + CAmount transparentBalance = 0; + CAmount sproutBalance = 0; + CAmount saplingBalance = 0; + for (const auto& t : spendableInputs.utxos) { + transparentBalance += t.Value(); + } + for (const auto& t : spendableInputs.sproutNoteEntries) { + sproutBalance += t.note.value(); + } + for (const auto& t : spendableInputs.saplingNoteEntries) { + saplingBalance += t.note.value(); + } + UniValue pools(UniValue::VOBJ); - pools.pushKV("transparent", 99999.99); - pools.pushKV("sprout", 99999.99); - pools.pushKV("sapling", 99999.99); - pools.pushKV("orchard", 99999.99); + auto renderBalance = [&](std::string poolName, CAmount balance) { + if (balance > 0) { + UniValue pool(UniValue::VOBJ); + pool.pushKV("valueZat", balance); + pools.pushKV(poolName, pool); + } + }; + renderBalance("transparent", transparentBalance); + renderBalance("sprout", sproutBalance); + renderBalance("sapling", saplingBalance); UniValue result(UniValue::VOBJ); result.pushKV("pools", pools);