wtxmgr/tx_test: extend TestRemoveUnminedTx to check balances

In this commit, we extend TestRemoveUnminedTx to also account for
checking the store's total balance (confirmed and unconfirmed). It
currently ensures that the UTXO state is correct, but as a sanity check,
we'll also ensure that balances are properly updated.

wtxmgr/tx_test.go

@@ -1324,59 +1324,142 @@ func TestRemoveUnminedTx(t *testing.T) {
 	}
 	defer teardown()
 
-	dbtx, err := db.BeginReadWriteTx()
+	// In order to reproduce real-world scenarios, we'll use a new database
+	// transaction for each interaction with the wallet.
+	//
+	// We'll start off the test by creating a new coinbase output at height
+	// 100 and inserting it into the store.
+	b100 := &BlockMeta{
+		Block: Block{Height: 100},
+		Time:  time.Now(),
+	}
+	initialBalance := int64(1e8)
+	cb := newCoinBase(initialBalance)
+	cbRec, err := NewTxRecordFromMsgTx(cb, b100.Time)
 	if err != nil {
 		t.Fatal(err)
 	}
-	defer dbtx.Commit()
+	commitDBTx(t, store, db, func(ns walletdb.ReadWriteBucket) {
-	ns := dbtx.ReadWriteBucket(namespaceKey)
+		if err := store.InsertTx(ns, cbRec, b100); err != nil {
+			t.Fatal(err)
+		}
+		err := store.AddCredit(ns, cbRec, b100, 0, false)
+		if err != nil {
+			t.Fatal(err)
+		}
+	})
 
-	// We'll start off by adding an unconfirmed transaction to the
+	// Determine the maturity height for the coinbase output created.
-	// transaction store.
+	coinbaseMaturity := int32(chaincfg.TestNet3Params.CoinbaseMaturity)
-	tx := TstRecvTx
+	maturityHeight := b100.Block.Height + coinbaseMaturity
-	txRec, err := wtxmgr.NewTxRecordFromMsgTx(tx.MsgTx(), time.Now())
+
+	// checkBalance is a helper function that compares the balance of the
+	// store with the expected value. The includeUnconfirmed boolean can be
+	// used to include the unconfirmed balance as a part of the total
+	// balance.
+	checkBalance := func(expectedBalance btcutil.Amount,
+		includeUnconfirmed bool) {
+
+		t.Helper()
+
+		minConfs := int32(1)
+		if includeUnconfirmed {
+			minConfs = 0
+		}
+
+		commitDBTx(t, store, db, func(ns walletdb.ReadWriteBucket) {
+			t.Helper()
+
+			b, err := store.Balance(ns, minConfs, maturityHeight)
+			if err != nil {
+				t.Fatalf("unable to retrieve balance: %v", err)
+			}
+			if b != expectedBalance {
+				t.Fatalf("expected balance of %d, got %d",
+					expectedBalance, b)
+			}
+		})
+	}
+
+	// Since we don't have any unconfirmed transactions within the store,
+	// the total balance reflecting confirmed and unconfirmed outputs should
+	// match the initial balance.
+	checkBalance(btcutil.Amount(initialBalance), false)
+	checkBalance(btcutil.Amount(initialBalance), true)
+
+	// Then, we'll create an unconfirmed spend for the coinbase output and
+	// insert it into the store.
+	b101 := &BlockMeta{
+		Block: Block{Height: 201},
+		Time:  time.Now(),
+	}
+	changeAmount := int64(4e7)
+	spendTx := spendOutput(&cbRec.Hash, 0, 5e7, changeAmount)
+	spendTxRec, err := NewTxRecordFromMsgTx(spendTx, b101.Time)
 	if err != nil {
-		t.Fatalf("unable to create unmined txns: %v", err)
+		t.Fatal(err)
-	}
-	if err := store.InsertTx(ns, txRec, nil); err != nil {
-		t.Fatalf("unable to insert transaction: %v", err)
 	}
+	commitDBTx(t, store, db, func(ns walletdb.ReadWriteBucket) {
+		if err := store.InsertTx(ns, spendTxRec, nil); err != nil {
+			t.Fatal(err)
+		}
+		err := store.AddCredit(ns, spendTxRec, nil, 1, true)
+		if err != nil {
+			t.Fatal(err)
+		}
+	})
 
-	// With the transaction inserted, ensure that it's reflected in the set
+	// With the unconfirmed spend inserted into the store, we'll query it
-	// of unmined transactions.
+	// for its unconfirmed tranasctions to ensure it was properly added.
-	unminedTxns, err := store.UnminedTxs(ns)
+	commitDBTx(t, store, db, func(ns walletdb.ReadWriteBucket) {
-	if err != nil {
+		unminedTxs, err := store.UnminedTxs(ns)
-		t.Fatalf("unable to query for unmined txns: %v", err)
+		if err != nil {
-	}
+			t.Fatalf("unable to query for unmined txs: %v", err)
-	if len(unminedTxns) != 1 {
+		}
-		t.Fatalf("expected 1 mined tx, instead got %v",
+		if len(unminedTxs) != 1 {
-			len(unminedTxns))
+			t.Fatalf("expected 1 mined tx, instead got %v",
-	}
+				len(unminedTxs))
-	unminedTxHash := unminedTxns[0].TxHash()
+		}
-	txHash := tx.MsgTx().TxHash()
+		unminedTxHash := unminedTxs[0].TxHash()
-	if !unminedTxHash.IsEqual(&txHash) {
+		spendTxHash := spendTx.TxHash()
-		t.Fatalf("mismatch tx hashes: expected %v, got %v",
+		if !unminedTxHash.IsEqual(&spendTxHash) {
-			tx.MsgTx().TxHash(), unminedTxHash)
+			t.Fatalf("mismatch tx hashes: expected %v, got %v",
-	}
+				spendTxHash, unminedTxHash)
+		}
+	})
 
-	// Next, we'll delete the unmined transaction in order to simulate an
+	// Now that an unconfirmed spend exists, there should no longer be any
-	// encountered conflict.
+	// confirmed balance. The total balance should now all be unconfirmed
-	if err := store.RemoveUnminedTx(ns, txRec); err != nil {
+	// and it should match the change amount of the unconfirmed spend
-		t.Fatalf("unable to remove unmined txns: %v", err)
+	// tranasction.
-	}
+	checkBalance(0, false)
+	checkBalance(btcutil.Amount(changeAmount), true)
 
-	// If we query again for the set of unconfirmed transactions, then we
+	// Now, we'll remove the unconfirmed spend tranaction from the store.
-	// should get an empty slice.
+	commitDBTx(t, store, db, func(ns walletdb.ReadWriteBucket) {
-	// With the transaction inserted, ensure that it's reflected in the set
+		if err := store.RemoveUnminedTx(ns, spendTxRec); err != nil {
-	// of unmined transactions.
+			t.Fatal(err)
-	unminedTxns, err = store.UnminedTxs(ns)
+		}
-	if err != nil {
+	})
-		t.Fatalf("unable to query for unmined txns: %v", err)
+
-	}
+	// We'll query the store one last time for its unconfirmed transactions
-	if len(unminedTxns) != 0 {
+	// to ensure the unconfirmed spend was properly removed above.
-		t.Fatalf("expected zero unmined txns, instead have %v",
+	commitDBTx(t, store, db, func(ns walletdb.ReadWriteBucket) {
-			len(unminedTxns))
+		unminedTxs, err := store.UnminedTxs(ns)
-	}
+		if err != nil {
+			t.Fatalf("unable to query for unmined txs: %v", err)
+		}
+		if len(unminedTxs) != 0 {
+			t.Fatalf("expected 0 mined txs, instead got %v",
+				len(unminedTxs))
+		}
+	})
+
+	// Finally, the total balance (including confirmed and unconfirmed)
+	// should once again match the initial balance, as the uncofirmed spend
+	// has already been removed.
+	checkBalance(btcutil.Amount(initialBalance), false)
+	checkBalance(btcutil.Amount(initialBalance), true)
 }
 
 // commitDBTx is a helper function that allows us to perform multiple operations