add test for getting balances

ubclaunchpad · Aug 12, 2017 · 7b02913 · 7b02913
1 parent f12ebbc
commit 7b02913
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/blockchain/wallet.go b/blockchain/wallet.go
@@ -220,7 +220,7 @@ func (w *Wallet) IsPending(txn *Transaction) (bool, int) {
 func (w *Wallet) GetEffectiveBalance() uint64 {
 	r := w.Balance
 	for _, t := range w.PendingTxns {
-		r -= t.Outputs[0].Amount
+		r -= t.GetTotalOutput()
 	}
 	return r
 }

diff --git a/blockchain/wallet_test.go b/blockchain/wallet_test.go
@@ -41,3 +41,13 @@ func TestDropAllPending(t *testing.T) {
 	result, _ = w.IsPending(txn)
 	assert.False(t, result)
 }
+
+func TestGetWalletBalances(t *testing.T) {
+	w := NewWallet()
+	txn := NewTestTransaction()
+	w.SetBalance(txn.GetTotalOutput())
+	w.SetAllPending([]*Transaction{txn})
+
+	assert.Equal(t, w.GetBalance(), txn.GetTotalOutput())
+	assert.Equal(t, w.GetEffectiveBalance(), uint64(0))
+}