[lbry] coin selection and balance no longer include stakes

rpc/api.proto

@@ -203,6 +203,7 @@ message FundTransactionRequest {
 	int32 required_confirmations = 3;
 	bool include_immature_coinbases = 4;
 	bool include_change_script = 5;
+	bool include_stakes = 6;
 }
 message FundTransactionResponse {
 	message PreviousOutput {

rpc/legacyrpc/methods.go

@@ -585,7 +585,7 @@ func getBalance(icmd interface{}, w *wallet.Wallet) (interface{}, error) {
 		accountName = *cmd.Account
 	}
 	if accountName == "*" {
-		balance, err = w.CalculateBalance(int32(*cmd.MinConf))
+		balance, _, err = w.CalculateBalance(int32(*cmd.MinConf))
 		if err != nil {
 			return nil, err
 		}
@@ -640,7 +640,7 @@ func getInfo(icmd interface{}, w *wallet.Wallet, chainClient *chain.RPCClient) (
 		return nil, err
 	}
 
-	bal, err := w.CalculateBalance(1)
+	bal, staked, err := w.CalculateBalance(1)
 	if err != nil {
 		return nil, err
 	}
@@ -649,6 +649,8 @@ func getInfo(icmd interface{}, w *wallet.Wallet, chainClient *chain.RPCClient) (
 	// to using the manager version.
 	info.WalletVersion = int32(waddrmgr.LatestMgrVersion)
 	info.Balance = bal.ToBTC()
+	_ = staked // TODO: add this to lbcd:
+	// info.Staked = staked.ToBTC()
 	info.PaytxFee = float64(txrules.DefaultRelayFeePerKb)
 	// We don't set the following since they don't make much sense in the
 	// wallet architecture:
@@ -834,7 +836,6 @@ func lookupKeyScope(kind *string) (waddrmgr.KeyScope, error) {
 	if kind == nil {
 		return waddrmgr.KeyScopeBIP0044, nil
 	}
-	// must be one of legacy / p2pkh or p2sh-p2wkh / p2sh-segwit, or p2wkh / bech32
 	switch strings.ToLower(*kind) {
 	case "legacy", "p2pkh": // could add "default" but it might be confused with the 1st parameter
 		return waddrmgr.KeyScopeBIP0044, nil

rpc/rpcserver/server.go

@@ -306,6 +306,7 @@ func (s *walletServer) FundTransaction(ctx context.Context, req *pb.FundTransact
 	policy := wallet.OutputSelectionPolicy{
 		Account:               req.Account,
 		RequiredConfirmations: req.RequiredConfirmations,
+		IncludeStakes:         req.IncludeStakes,
 	}
 	unspentOutputs, err := s.wallet.UnspentOutputs(policy)
 	if err != nil {

rpc/walletrpc/api.pb.go

@@ -892,6 +892,7 @@ type FundTransactionRequest struct {
 	RequiredConfirmations    int32  `protobuf:"varint,3,opt,name=required_confirmations,json=requiredConfirmations" json:"required_confirmations,omitempty"`
 	IncludeImmatureCoinbases bool   `protobuf:"varint,4,opt,name=include_immature_coinbases,json=includeImmatureCoinbases" json:"include_immature_coinbases,omitempty"`
 	IncludeChangeScript      bool   `protobuf:"varint,5,opt,name=include_change_script,json=includeChangeScript" json:"include_change_script,omitempty"`
+	IncludeStakes            bool   `protobuf:"varint,6,opt,name=include_stakes,json=includeStakes" json:"include_stakes,omitempty"`
 }
 
 func (m *FundTransactionRequest) Reset()                    { *m = FundTransactionRequest{} }

wallet/common.go

@@ -9,6 +9,7 @@ import (
 	"time"
 
 	"github.com/lbryio/lbcd/chaincfg/chainhash"
+	"github.com/lbryio/lbcd/txscript"
 	"github.com/lbryio/lbcd/wire"
 	btcutil "github.com/lbryio/lbcutil"
 )
@@ -41,8 +42,18 @@ type OutputKind byte
 const (
 	OutputKindNormal OutputKind = iota
 	OutputKindCoinbase
+	OutputKindStake
 )
 
+func isStake(pkScript []byte) bool {
+	if len(pkScript) > 0 &&
+		(pkScript[0] == txscript.OP_CLAIMNAME || pkScript[0] == txscript.OP_SUPPORTCLAIM ||
+			pkScript[0] == txscript.OP_UPDATECLAIM) {
+		return true
+	}
+	return false
+}
+
 // TransactionOutput describes an output that was or is at least partially
 // controlled by the wallet.  Depending on context, this could refer to an
 // unspent output, or a spent one.

wallet/createtx.go

@@ -316,6 +316,11 @@ func (w *Wallet) findEligibleOutputs(dbtx walletdb.ReadTx,
 		if err != nil || len(addrs) != 1 {
 			continue
 		}
+
+		if isStake(output.PkScript) {
+			continue
+		}
+
 		scopedMgr, addrAcct, err := w.Manager.AddrAccount(addrmgrNs, addrs[0])
 		if err != nil {
 			continue

wallet/utxos.go

@@ -29,6 +29,7 @@ var (
 type OutputSelectionPolicy struct {
 	Account               uint32
 	RequiredConfirmations int32
+	IncludeStakes         bool
 }
 
 func (p *OutputSelectionPolicy) meetsRequiredConfs(txHeight, curHeight int32) bool {
@@ -84,6 +85,13 @@ func (w *Wallet) UnspentOutputs(policy OutputSelectionPolicy) ([]*TransactionOut
 				outputSource = OutputKindCoinbase
 			}
 
+			if isStake(output.PkScript) {
+				if !policy.IncludeStakes {
+					continue
+				}
+				outputSource = OutputKindStake
+			}
+
 			result := &TransactionOutput{
 				OutPoint: output.OutPoint,
 				Output: wire.TxOut{

wallet/wallet.go

@@ -1462,16 +1462,17 @@ func (w *Wallet) AccountAddresses(account uint32) (addrs []btcutil.Address, err
 // a UTXO must be in a block.  If confirmations is 1 or greater,
 // the balance will be calculated based on how many how many blocks
 // include a UTXO.
-func (w *Wallet) CalculateBalance(confirms int32) (btcutil.Amount, error) {
+func (w *Wallet) CalculateBalance(confirms int32) (btcutil.Amount, btcutil.Amount, error) {
 	var balance btcutil.Amount
+	var staked btcutil.Amount
 	err := walletdb.View(w.db, func(tx walletdb.ReadTx) error {
 		txmgrNs := tx.ReadBucket(wtxmgrNamespaceKey)
 		var err error
 		blk := w.Manager.SyncedTo()
-		balance, err = w.TxStore.Balance(txmgrNs, confirms, blk.Height)
+		balance, staked, err = w.TxStore.Balance(txmgrNs, confirms, blk.Height)
 		return err
 	})
-	return balance, err
+	return balance, staked, err
 }
 
 // Balances records total, spendable (by policy), and immature coinbase
@@ -1511,7 +1512,7 @@ func (w *Wallet) CalculateAccountBalances(account uint32, confirms int32) (Balan
 			if err == nil && len(addrs) > 0 {
 				_, outputAcct, err = w.Manager.AddrAccount(addrmgrNs, addrs[0])
 			}
-			if err != nil || outputAcct != account {
+			if err != nil || outputAcct != account || isStake(output.PkScript) {
 				continue
 			}
 
@@ -2478,6 +2479,9 @@ func (w *Wallet) AccountBalances(scope waddrmgr.KeyScope,
 			if err != nil || len(addrs) == 0 {
 				continue
 			}
+			if isStake(output.PkScript) {
+				continue
+			}
 			outputAcct, err := manager.AddrAccount(addrmgrNs, addrs[0])
 			if err != nil {
 				continue
@@ -2651,7 +2655,8 @@ func (w *Wallet) ListUnspent(minconf, maxconf int32,
 				ScriptPubKey:  hex.EncodeToString(output.PkScript),
 				Amount:        output.Amount.ToBTC(),
 				Confirmations: int64(confs),
-				Spendable:     spendable,
+				Spendable:     spendable, // presently false for stakes
+				// TODO: add an IsStake flag here to lbcd
 			}
 
 			// BUG: this should be a JSON array so that all

wtxmgr/db.go

@@ -543,9 +543,7 @@ func keyCredit(txHash *chainhash.Hash, index uint32, block *Block) []byte {
 func valueUnspentCredit(cred *credit) []byte {
 	v := make([]byte, 9)
 	byteOrder.PutUint64(v, uint64(cred.amount))
-	if cred.change {
+	v[8] = cred.flags
-		v[8] |= 1 << 1
-	}
 	return v
 }
 
@@ -598,13 +596,13 @@ func fetchRawCreditAmountSpent(v []byte) (btcutil.Amount, bool, error) {
 
 // fetchRawCreditAmountChange returns the amount of the credit and whether the
 // credit is marked as change.
-func fetchRawCreditAmountChange(v []byte) (btcutil.Amount, bool, error) {
+func fetchRawCreditAmountChange(v []byte) (btcutil.Amount, byte, error) {
 	if len(v) < 9 {
 		str := fmt.Sprintf("%s: short read (expected %d bytes, read %d)",
 			bucketCredits, 9, len(v))
-		return 0, false, storeError(ErrData, str, nil)
+		return 0, 0, storeError(ErrData, str, nil)
 	}
-	return btcutil.Amount(byteOrder.Uint64(v)), v[8]&(1<<1) != 0, nil
+	return btcutil.Amount(byteOrder.Uint64(v)), v[8], nil
 }
 
 // fetchRawCreditUnspentValue returns the unspent value for a raw credit key.
@@ -1037,12 +1035,10 @@ func deleteRawUnmined(ns walletdb.ReadWriteBucket, k []byte) error {
 //   [8]     Flags (1 byte)
 //             0x02: Change
 
-func valueUnminedCredit(amount btcutil.Amount, change bool) []byte {
+func valueUnminedCredit(amount btcutil.Amount, flags byte) []byte {
 	v := make([]byte, 9)
 	byteOrder.PutUint64(v, uint64(amount))
-	if change {
+	v[8] = flags
-		v[8] = 1 << 1
-	}
 	return v
 }
 
@@ -1071,14 +1067,13 @@ func fetchRawUnminedCreditAmount(v []byte) (btcutil.Amount, error) {
 	return btcutil.Amount(byteOrder.Uint64(v)), nil
 }
 
-func fetchRawUnminedCreditAmountChange(v []byte) (btcutil.Amount, bool, error) {
+func fetchRawUnminedCreditAmountChange(v []byte) (btcutil.Amount, byte, error) {
 	if len(v) < 9 {
 		str := "short unmined credit value"
-		return 0, false, storeError(ErrData, str, nil)
+		return 0, 0, storeError(ErrData, str, nil)
 	}
 	amt := btcutil.Amount(byteOrder.Uint64(v))
-	change := v[8]&(1<<1) != 0
+	return amt, v[8], nil
-	return amt, change, nil
 }
 
 func existsRawUnminedCredit(ns walletdb.ReadBucket, k []byte) []byte {
@@ -1146,14 +1141,14 @@ func (it *unminedCreditIterator) readElem() error {
 	if err != nil {
 		return err
 	}
-	amount, change, err := fetchRawUnminedCreditAmountChange(it.cv)
+	amount, flags, err := fetchRawUnminedCreditAmountChange(it.cv)
 	if err != nil {
 		return err
 	}
 
 	it.elem.Index = index
 	it.elem.Amount = amount
-	it.elem.Change = change
+	it.elem.Change = (flags & ChangeFlag) > 0
 	// Spent intentionally not set
 
 	return nil

wtxmgr/example_test.go

@@ -232,7 +232,7 @@ func Example_basicUsage() {
 	}
 
 	// Print the one confirmation balance.
-	bal, err := s.Balance(b, 1, 100)
+	bal, _, err := s.Balance(b, 1, 100)
 	if err != nil {
 		fmt.Println(err)
 		return

wtxmgr/tx.go

@@ -15,6 +15,7 @@ import (
 	"github.com/lbryio/lbcd/blockchain"
 	"github.com/lbryio/lbcd/chaincfg"
 	"github.com/lbryio/lbcd/chaincfg/chainhash"
+	"github.com/lbryio/lbcd/txscript"
 	"github.com/lbryio/lbcd/wire"
 	btcutil "github.com/lbryio/lbcutil"
 	"github.com/lbryio/lbcwallet/walletdb"
@@ -24,6 +25,9 @@ import (
 const (
 	// TxLabelLimit is the length limit we impose on transaction labels.
 	TxLabelLimit = 500
+
+	ChangeFlag = 2
+	StakeFlag  = 4
 )
 
 var (
@@ -114,7 +118,7 @@ type credit struct {
 	outPoint wire.OutPoint
 	block    Block
 	amount   btcutil.Amount
-	change   bool
+	flags    byte
 	spentBy  indexedIncidence // Index == ^uint32(0) if unspent
 }
 
@@ -304,14 +308,14 @@ func (s *Store) updateMinedBalance(ns walletdb.ReadWriteBucket, rec *TxRecord,
 		if err != nil {
 			return err
 		}
-		amount, change, err := fetchRawUnminedCreditAmountChange(it.cv)
+		amount, flags, err := fetchRawUnminedCreditAmountChange(it.cv)
 		if err != nil {
 			return err
 		}
 
 		cred.outPoint.Index = index
 		cred.amount = amount
-		cred.change = change
+		cred.flags = flags
 
 		if err := putUnspentCredit(ns, &cred); err != nil {
 			return err
@@ -494,6 +498,11 @@ func (s *Store) AddCredit(ns walletdb.ReadWriteBucket, rec *TxRecord, block *Blo
 // bool return specifies whether the unspent output is newly added (true) or a
 // duplicate (false).
 func (s *Store) addCredit(ns walletdb.ReadWriteBucket, rec *TxRecord, block *BlockMeta, index uint32, change bool) (bool, error) {
+	flags := isStake(rec.MsgTx.TxOut[index])
+	if change {
+		flags |= ChangeFlag
+	}
+
 	if block == nil {
 		// If the outpoint that we should mark as credit already exists
 		// within the store, either as unconfirmed or confirmed, then we
@@ -507,7 +516,7 @@ func (s *Store) addCredit(ns walletdb.ReadWriteBucket, rec *TxRecord, block *Blo
 				rec.Hash.String())
 			return false, nil
 		}
-		v := valueUnminedCredit(btcutil.Amount(rec.MsgTx.TxOut[index].Value), change)
+		v := valueUnminedCredit(btcutil.Amount(rec.MsgTx.TxOut[index].Value), flags)
 		return true, putRawUnminedCredit(ns, k, v)
 	}
 
@@ -527,7 +536,7 @@ func (s *Store) addCredit(ns walletdb.ReadWriteBucket, rec *TxRecord, block *Blo
 		},
 		block:   block.Block,
 		amount:  txOutAmt,
-		change:  change,
+		flags:   flags,
 		spentBy: indexedIncidence{index: ^uint32(0)},
 	}
 	v = valueUnspentCredit(&cred)
@@ -548,6 +557,15 @@ func (s *Store) addCredit(ns walletdb.ReadWriteBucket, rec *TxRecord, block *Blo
 	return true, putUnspent(ns, &cred.outPoint, &block.Block)
 }
 
+func isStake(out *wire.TxOut) byte {
+	if len(out.PkScript) > 0 &&
+		(out.PkScript[0] == txscript.OP_CLAIMNAME || out.PkScript[0] == txscript.OP_SUPPORTCLAIM ||
+			out.PkScript[0] == txscript.OP_UPDATECLAIM) {
+		return StakeFlag
+	}
+	return 0
+}
+
 // Rollback removes all blocks at height onwards, moving any transactions within
 // each block to the unconfirmed pool.
 func (s *Store) Rollback(ns walletdb.ReadWriteBucket, height int32) error {
@@ -710,12 +728,12 @@ func (s *Store) rollback(ns walletdb.ReadWriteBucket, height int32) error {
 					continue
 				}
 
-				amt, change, err := fetchRawCreditAmountChange(v)
+				amt, flags, err := fetchRawCreditAmountChange(v)
 				if err != nil {
 					return err
 				}
 				outPointKey := canonicalOutPoint(&rec.Hash, uint32(i))
-				unminedCredVal := valueUnminedCredit(amt, change)
+				unminedCredVal := valueUnminedCredit(amt, flags)
 				err = putRawUnminedCredit(ns, outPointKey, unminedCredVal)
 				if err != nil {
 					return err
@@ -918,14 +936,14 @@ func (s *Store) UnspentOutputs(ns walletdb.ReadBucket) ([]Credit, error) {
 //
 // Balance may return unexpected results if syncHeight is lower than the block
 // height of the most recent mined transaction in the store.
-func (s *Store) Balance(ns walletdb.ReadBucket, minConf int32, syncHeight int32) (btcutil.Amount, error) {
+func (s *Store) Balance(ns walletdb.ReadBucket, minConf int32, syncHeight int32) (btcutil.Amount, btcutil.Amount, error) {
 	bal, err := fetchMinedBalance(ns)
 	if err != nil {
-		return 0, err
+		return 0, 0, err
 	}
 
-	// Subtract the balance for each credit that is spent by an unmined
+	// Subtract the balance for each credit that is spent by an unmined transaction or moved to stake.
-	// transaction.
+	var staked btcutil.Amount
 	var op wire.OutPoint
 	var block Block
 	err = ns.NestedReadBucket(bucketUnspent).ForEach(func(k, v []byte) error {
@@ -938,14 +956,15 @@ func (s *Store) Balance(ns walletdb.ReadBucket, minConf int32, syncHeight int32)
 			return err
 		}
 
+		_, c := existsCredit(ns, &op.Hash, op.Index, &block)
+		amt, flags, err := fetchRawCreditAmountChange(c)
+		if err != nil {
+			return err
+		}
+
 		// Subtract the output's amount if it's locked.
 		_, _, isLocked := isLockedOutput(ns, op, s.clock.Now())
 		if isLocked {
-			_, v := existsCredit(ns, &op.Hash, op.Index, &block)
-			amt, err := fetchRawCreditAmount(v)
-			if err != nil {
-				return err
-			}
 			bal -= amt
 
 			// To prevent decrementing the balance twice if the
@@ -954,22 +973,20 @@ func (s *Store) Balance(ns walletdb.ReadBucket, minConf int32, syncHeight int32)
 		}
 
 		if existsRawUnminedInput(ns, k) != nil {
-			_, v := existsCredit(ns, &op.Hash, op.Index, &block)
-			amt, err := fetchRawCreditAmount(v)
-			if err != nil {
-				return err
-			}
 			bal -= amt
+		} else if (flags & StakeFlag) > 0 {
+			bal -= amt
+			staked += amt
 		}
 
 		return nil
 	})
 	if err != nil {
 		if _, ok := err.(Error); ok {
-			return 0, err
+			return 0, 0, err
 		}
 		str := "failed iterating unspent outputs"
-		return 0, storeError(ErrDatabase, str, err)
+		return 0, 0, storeError(ErrDatabase, str, err)
 	}
 
 	// Decrement the balance for any unspent credit with less than
@@ -992,7 +1009,7 @@ func (s *Store) Balance(ns walletdb.ReadBucket, minConf int32, syncHeight int32)
 			txHash := &block.transactions[i]
 			rec, err := fetchTxRecord(ns, txHash, &block.Block)
 			if err != nil {
-				return 0, err
+				return 0, 0, err
 			}
 			numOuts := uint32(len(rec.MsgTx.TxOut))
 			for i := uint32(0); i < numOuts; i++ {
@@ -1017,7 +1034,7 @@ func (s *Store) Balance(ns walletdb.ReadBucket, minConf int32, syncHeight int32)
 				}
 				amt, spent, err := fetchRawCreditAmountSpent(v)
 				if err != nil {
-					return 0, err
+					return 0, 0, err
 				}
 				if spent {
 					continue
@@ -1031,7 +1048,7 @@ func (s *Store) Balance(ns walletdb.ReadBucket, minConf int32, syncHeight int32)
 		}
 	}
 	if blockIt.err != nil {
-		return 0, blockIt.err
+		return 0, 0, blockIt.err
 	}
 
 	// If unmined outputs are included, increment the balance for each
@@ -1064,14 +1081,14 @@ func (s *Store) Balance(ns walletdb.ReadBucket, minConf int32, syncHeight int32)
 		})
 		if err != nil {
 			if _, ok := err.(Error); ok {
-				return 0, err
+				return 0, 0, err
 			}
 			str := "failed to iterate over unmined credits bucket"
-			return 0, storeError(ErrDatabase, str, err)
+			return 0, 0, storeError(ErrDatabase, str, err)
 		}
 	}
 
-	return bal, nil
+	return bal, staked, nil
 }
 
 // PutTxLabel validates transaction labels and writes them to disk if they

wtxmgr/tx_test.go

@@ -526,14 +526,14 @@ func TestInsertsCreditsDebitsRollbacks(t *testing.T) {
 				t.Fatalf("%s: got error: %v", test.name, err)
 			}
 			s = tmpStore
-			bal, err := s.Balance(ns, 1, TstRecvCurrentHeight)
+			bal, _, err := s.Balance(ns, 1, TstRecvCurrentHeight)
 			if err != nil {
 				t.Fatalf("%s: Confirmed Balance failed: %v", test.name, err)
 			}
 			if bal != test.bal {
 				t.Fatalf("%s: balance mismatch: expected: %d, got: %d", test.name, test.bal, bal)
 			}
-			unc, err := s.Balance(ns, 0, TstRecvCurrentHeight)
+			unc, _, err := s.Balance(ns, 0, TstRecvCurrentHeight)
 			if err != nil {
 				t.Fatalf("%s: Unconfirmed Balance failed: %v", test.name, err)
 			}
@@ -626,7 +626,7 @@ func TestFindingSpentCredits(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	bal, err := s.Balance(ns, 1, TstSignedTxBlockDetails.Height)
+	bal, _, err := s.Balance(ns, 1, TstSignedTxBlockDetails.Height)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -809,7 +809,7 @@ func TestCoinbases(t *testing.T) {
 		},
 	}
 	for i, tst := range balTests {
-		bal, err := s.Balance(ns, tst.minConf, tst.height)
+		bal, _, err := s.Balance(ns, tst.minConf, tst.height)
 		if err != nil {
 			t.Fatalf("Balance test %d: Store.Balance failed: %v", i, err)
 		}
@@ -890,7 +890,7 @@ func TestCoinbases(t *testing.T) {
 	}
 	balTestsBeforeMaturity := balTests
 	for i, tst := range balTests {
-		bal, err := s.Balance(ns, tst.minConf, tst.height)
+		bal, _, err := s.Balance(ns, tst.minConf, tst.height)
 		if err != nil {
 			t.Fatalf("Balance test %d: Store.Balance failed: %v", i, err)
 		}
@@ -973,7 +973,7 @@ func TestCoinbases(t *testing.T) {
 		},
 	}
 	for i, tst := range balTests {
-		bal, err := s.Balance(ns, tst.minConf, tst.height)
+		bal, _, err := s.Balance(ns, tst.minConf, tst.height)
 		if err != nil {
 			t.Fatalf("Balance test %d: Store.Balance failed: %v", i, err)
 		}
@@ -1006,7 +1006,7 @@ func TestCoinbases(t *testing.T) {
 		t.Fatal(err)
 	}
 	for i, tst := range balTests {
-		bal, err := s.Balance(ns, tst.minConf, tst.height)
+		bal, _, err := s.Balance(ns, tst.minConf, tst.height)
 		if err != nil {
 			t.Fatalf("Balance test %d: Store.Balance failed: %v", i, err)
 		}
@@ -1026,7 +1026,7 @@ func TestCoinbases(t *testing.T) {
 	}
 	balTests = balTestsBeforeMaturity
 	for i, tst := range balTests {
-		bal, err := s.Balance(ns, tst.minConf, tst.height)
+		bal, _, err := s.Balance(ns, tst.minConf, tst.height)
 		if err != nil {
 			t.Fatalf("Balance test %d: Store.Balance failed: %v", i, err)
 		}
@@ -1072,7 +1072,7 @@ func TestCoinbases(t *testing.T) {
 		},
 	}
 	for i, tst := range balTests {
-		bal, err := s.Balance(ns, tst.minConf, tst.height)
+		bal, _, err := s.Balance(ns, tst.minConf, tst.height)
 		if err != nil {
 			t.Fatalf("Balance test %d: Store.Balance failed: %v", i, err)
 		}
@@ -1247,7 +1247,7 @@ func TestMoveMultipleToSameBlock(t *testing.T) {
 		},
 	}
 	for i, tst := range balTests {
-		bal, err := s.Balance(ns, tst.minConf, tst.height)
+		bal, _, err := s.Balance(ns, tst.minConf, tst.height)
 		if err != nil {
 			t.Fatalf("Balance test %d: Store.Balance failed: %v", i, err)
 		}
@@ -1392,7 +1392,7 @@ func TestRemoveUnminedTx(t *testing.T) {
 		commitDBTx(t, store, db, func(ns walletdb.ReadWriteBucket) {
 			t.Helper()
 
-			b, err := store.Balance(ns, minConfs, maturityHeight)
+			b, _, err := store.Balance(ns, minConfs, maturityHeight)
 			if err != nil {
 				t.Fatalf("unable to retrieve balance: %v", err)
 			}
@@ -2389,7 +2389,7 @@ func assertBalance(t *testing.T, s *Store, ns walletdb.ReadWriteBucket,
 	if confirmed {
 		minConf = 1
 	}
-	balance, err := s.Balance(ns, minConf, blockHeight)
+	balance, _, err := s.Balance(ns, minConf, blockHeight)
 	if err != nil {
 		t.Fatal(err)
 	}