From 05bf4a645959e8cafef6581fa6b4f9add538cb55 Mon Sep 17 00:00:00 2001
From: Aditya Kulkarni <adityapk@gmail.com>
Date: Tue, 29 Oct 2019 10:13:34 -0700
Subject: [PATCH] Add check for balance/zbalance

---
 lib/src/lightwallet/tests.rs | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/lib/src/lightwallet/tests.rs b/lib/src/lightwallet/tests.rs
index 2997694..69fdbae 100644
--- a/lib/src/lightwallet/tests.rs
+++ b/lib/src/lightwallet/tests.rs
@@ -706,6 +706,12 @@ fn test_z_spend_to_z() {
     let branch_id = u32::from_str_radix("2bb40e60", 16).unwrap();
     let (ss, so) = get_sapling_params().unwrap();
 
+    // Make sure that the balance exists 
+    {
+        assert_eq!(wallet.zbalance(None), AMOUNT1);
+        assert_eq!(wallet.verified_zbalance(None), AMOUNT1);
+    }
+
     // Create a tx and send to address
     let raw_tx = wallet.send_to_address(branch_id, &ss, &so,
                             vec![(&ext_address, AMOUNT_SENT, Some(outgoing_memo.clone()))]).unwrap();
@@ -736,6 +742,12 @@ fn test_z_spend_to_z() {
         assert_eq!(mem[&sent_txid].outgoing_metadata[0].memo.to_utf8().unwrap().unwrap(), outgoing_memo);
     }
 
+    {
+        // The wallet should deduct this from the balance and verified balance
+        assert_eq!(wallet.zbalance(None), 0);
+        assert_eq!(wallet.verified_zbalance(None), 0);
+    }
+
     let mut cb3 = FakeCompactBlock::new(2, block_hash);
     cb3.add_tx(&sent_tx);
     wallet.scan_block(&cb3.as_bytes()).unwrap();