Add TestState::{get_balance, get_balance_at} convenience methods.

Signed-off-by: Daira Emma Hopwood <[email protected]>

zcash_client_sqlite/src/testing.rs

@@ -57,7 +57,10 @@ use zcash_primitives::{
 use crate::{
     chain::init::init_cache_database,
     error::SqliteClientError,
-    wallet::{commitment_tree, init::init_wallet_db, sapling::tests::test_prover},
+    wallet::{
+        commitment_tree, get_balance, get_balance_at, init::init_wallet_db,
+        sapling::tests::test_prover,
+    },
     AccountId, ReceivedNoteId, WalletDb,
 };
 
@@ -323,6 +326,23 @@ impl<Cache> TestState<Cache> {
             .and_then(|(_, usk, _)| usk.to_unified_full_viewing_key().sapling().cloned())
     }
 
+    /// Returns the balance for the account, including all mined unspent notes
+    /// that we know about.
+    pub(crate) fn get_balance(&self, account: AccountId) -> Result<Amount, SqliteClientError> {
+        get_balance(&self.wallet().conn, account)
+    }
+
+    /// Returns the verified balance for the account at the specified height,
+    /// This may be used to obtain a balance that ignores notes that have been
+    /// received so recently that they are not yet deemed spendable.
+    pub(crate) fn get_balance_at(
+        &self,
+        account: AccountId,
+        height: BlockHeight,
+    ) -> Result<Amount, SqliteClientError> {
+        get_balance_at(&self.wallet().conn, account, height)
+    }
+
     /// Invokes [`create_spend_to_address`] with the given arguments.
     #[allow(deprecated)]
     #[allow(clippy::type_complexity)]

zcash_client_sqlite/src/wallet/sapling.rs

@@ -440,7 +440,7 @@ pub(crate) mod tests {
     use crate::{
         error::SqliteClientError,
         testing::{AddressType, BlockCache, TestBuilder, TestState},
-        wallet::{commitment_tree, get_balance, get_balance_at},
+        wallet::commitment_tree,
         AccountId, NoteId, ReceivedNoteId,
     };
 
@@ -479,12 +479,10 @@ pub(crate) mod tests {
             .get_target_and_anchor_heights(NonZeroU32::new(1).unwrap())
             .unwrap()
             .unwrap();
+        assert_eq!(st.get_balance(AccountId::from(0)).unwrap(), value);
         assert_eq!(
-            get_balance(&st.wallet().conn, AccountId::from(0)).unwrap(),
-            value
-        );
-        assert_eq!(
-            get_balance_at(&st.wallet().conn, AccountId::from(0), anchor_height).unwrap(),
+            st.get_balance_at(AccountId::from(0), anchor_height)
+                .unwrap(),
             value
         );
 
@@ -646,10 +644,7 @@ pub(crate) mod tests {
         let to = dfvk.default_address().1.into();
 
         // Account balance should be zero
-        assert_eq!(
-            get_balance(&st.wallet().conn, AccountId::from(0)).unwrap(),
-            Amount::zero()
-        );
+        assert_eq!(st.get_balance(AccountId::from(0)).unwrap(), Amount::zero());
 
         // We cannot do anything if we aren't synchronised
         assert_matches!(
@@ -685,12 +680,10 @@ pub(crate) mod tests {
             .get_target_and_anchor_heights(NonZeroU32::new(10).unwrap())
             .unwrap()
             .unwrap();
+        assert_eq!(st.get_balance(AccountId::from(0)).unwrap(), value);
         assert_eq!(
-            get_balance(&st.wallet().conn, AccountId::from(0)).unwrap(),
-            value
-        );
-        assert_eq!(
-            get_balance_at(&st.wallet().conn, AccountId::from(0), anchor_height).unwrap(),
+            st.get_balance_at(AccountId::from(0), anchor_height)
+                .unwrap(),
             value
         );
 
@@ -705,11 +698,12 @@ pub(crate) mod tests {
             .unwrap()
             .unwrap();
         assert_eq!(
-            get_balance(&st.wallet().conn, AccountId::from(0)).unwrap(),
+            st.get_balance(AccountId::from(0)).unwrap(),
             (value + value).unwrap()
         );
         assert_eq!(
-            get_balance_at(&st.wallet().conn, AccountId::from(0), anchor_height2).unwrap(),
+            st.get_balance_at(AccountId::from(0), anchor_height2)
+                .unwrap(),
             value
         );
 
@@ -789,10 +783,7 @@ pub(crate) mod tests {
         let value = Amount::from_u64(50000).unwrap();
         let (h1, _, _) = st.generate_next_block(&dfvk, AddressType::DefaultExternal, value);
         st.scan_cached_blocks(h1, 1);
-        assert_eq!(
-            get_balance(&st.wallet().conn, AccountId::from(0)).unwrap(),
-            value
-        );
+        assert_eq!(st.get_balance(AccountId::from(0)).unwrap(), value);
 
         // Send some of the funds to another address
         let extsk2 = ExtendedSpendingKey::master(&[]);
@@ -887,10 +878,7 @@ pub(crate) mod tests {
         let value = Amount::from_u64(50000).unwrap();
         let (h1, _, _) = st.generate_next_block(&dfvk, AddressType::DefaultExternal, value);
         st.scan_cached_blocks(h1, 1);
-        assert_eq!(
-            get_balance(&st.wallet().conn, AccountId::from(0)).unwrap(),
-            value
-        );
+        assert_eq!(st.get_balance(AccountId::from(0)).unwrap(), value);
 
         let extsk2 = ExtendedSpendingKey::master(&[]);
         let addr2 = extsk2.default_address().1;
@@ -994,12 +982,10 @@ pub(crate) mod tests {
             .get_target_and_anchor_heights(NonZeroU32::new(1).unwrap())
             .unwrap()
             .unwrap();
+        assert_eq!(st.get_balance(AccountId::from(0)).unwrap(), value);
         assert_eq!(
-            get_balance(&st.wallet().conn, AccountId::from(0)).unwrap(),
-            value
-        );
-        assert_eq!(
-            get_balance_at(&st.wallet().conn, AccountId::from(0), anchor_height).unwrap(),
+            st.get_balance_at(AccountId::from(0), anchor_height)
+                .unwrap(),
             value
         );
 
@@ -1037,12 +1023,10 @@ pub(crate) mod tests {
             .get_target_and_anchor_heights(NonZeroU32::new(10).unwrap())
             .unwrap()
             .unwrap();
+        assert_eq!(st.get_balance(AccountId::from(0)).unwrap(), value);
         assert_eq!(
-            get_balance(&st.wallet().conn, AccountId::from(0)).unwrap(),
-            value
-        );
-        assert_eq!(
-            get_balance_at(&st.wallet().conn, AccountId::from(0), anchor_height).unwrap(),
+            st.get_balance_at(AccountId::from(0), anchor_height)
+                .unwrap(),
             value
         );
 
@@ -1094,12 +1078,10 @@ pub(crate) mod tests {
             .get_target_and_anchor_heights(NonZeroU32::new(1).unwrap())
             .unwrap()
             .unwrap();
+        assert_eq!(st.get_balance(AccountId::from(0)).unwrap(), total);
         assert_eq!(
-            get_balance(&st.wallet().conn, AccountId::from(0)).unwrap(),
-            total
-        );
-        assert_eq!(
-            get_balance_at(&st.wallet().conn, AccountId::from(0), anchor_height).unwrap(),
+            st.get_balance_at(AccountId::from(0), anchor_height)
+                .unwrap(),
             total
         );