Add tests for view only balance

full-service/src/service/balance.rs

@@ -403,12 +403,20 @@ where
 mod tests {
     use super::*;
     use crate::{
-        service::{account::AccountService, address::AddressService},
+        service::{
+            account::AccountService, address::AddressService,
+            view_only_account::ViewOnlyAccountService,
+        },
         test_utils::{get_test_ledger, manually_sync_account, setup_wallet_service, MOB},
         util::b58::b58_encode_public_address,
     };
-    use mc_account_keys::{AccountKey, PublicAddress, RootEntropy, RootIdentity};
+    use mc_account_keys::{
+        AccountKey, PublicAddress, RootEntropy, RootIdentity, CHANGE_SUBADDRESS_INDEX,
+        DEFAULT_SUBADDRESS_INDEX,
+    };
     use mc_common::logger::{test_with_logger, Logger};
+    use mc_crypto_keys::{RistrettoPrivate, RistrettoPublic};
+    use mc_transaction_core::{encrypted_fog_hint::EncryptedFogHint, tx::TxOut};
     use mc_util_from_random::FromRandom;
     use rand::{rngs::StdRng, SeedableRng};
 
@@ -510,4 +518,130 @@ mod tests {
             Err(e) => panic!("Unexpected error {:?}", e),
         }
     }
+
+    // The balance for an address should be accurate.
+    #[test_with_logger]
+    fn test_view_only_balance(logger: Logger) {
+        // setup view only account
+        let mut rng: StdRng = SeedableRng::from_seed([20u8; 32]);
+        let known_recipients: Vec<PublicAddress> = Vec::new();
+        let current_block_height = 12; //index 11
+        let ledger_db = get_test_ledger(
+            5,
+            &known_recipients,
+            current_block_height as usize,
+            &mut rng,
+        );
+        let service = setup_wallet_service(ledger_db.clone(), logger.clone());
+        let conn = service.wallet_db.get_conn().unwrap();
+
+        let view_private_key = RistrettoPrivate::from_random(&mut rng);
+        let spend_private_key = RistrettoPrivate::from_random(&mut rng);
+
+        let name = "testing";
+
+        let account_key = AccountKey::new(&spend_private_key, &view_private_key);
+        let account_id = AccountID::from(&account_key);
+        let main_public_address = account_key.default_subaddress();
+        let change_public_address = account_key.change_subaddress();
+        let mut subaddresses: Vec<(String, u64, String, RistrettoPublic)> = Vec::new();
+        subaddresses.push((
+            b58_encode_public_address(&main_public_address).unwrap(),
+            DEFAULT_SUBADDRESS_INDEX,
+            "Main".to_string(),
+            *main_public_address.spend_public_key(),
+        ));
+        subaddresses.push((
+            b58_encode_public_address(&change_public_address).unwrap(),
+            CHANGE_SUBADDRESS_INDEX,
+            "Change".to_string(),
+            *change_public_address.spend_public_key(),
+        ));
+
+        service
+            .import_view_only_account(
+                &account_id.to_string(),
+                &view_private_key,
+                DEFAULT_SUBADDRESS_INDEX,
+                CHANGE_SUBADDRESS_INDEX,
+                2,
+                name.clone(),
+                subaddresses,
+            )
+            .unwrap();
+
+        // add funds to account
+        for _ in 0..2 {
+            let value = 420 * MOB;
+            let tx_private_key = RistrettoPrivate::from_random(&mut rng);
+            let hint = EncryptedFogHint::fake_onetime_hint(&mut rng);
+            let fake_tx_out =
+                TxOut::new(value as u64, &main_public_address, &tx_private_key, hint).unwrap();
+            ViewOnlyTxo::create(
+                fake_tx_out.clone(),
+                value,
+                Some(DEFAULT_SUBADDRESS_INDEX),
+                Some(current_block_height),
+                &account_id.to_string(),
+                &conn,
+            )
+            .unwrap();
+        }
+
+        // test balance for account
+        let balance: Balance = service
+            .get_balance_for_view_only_account(&account_id.to_string())
+            .unwrap();
+        assert_eq!(balance.unspent as u64, 840 * MOB);
+        // view only accounts have no spendable MOB
+        assert_eq!(balance.max_spendable, 0);
+        assert_eq!(balance.spent, 0);
+        assert_eq!(balance.pending, 0);
+        assert_eq!(balance.secreted, 0);
+        assert_eq!(balance.orphaned, 0);
+
+        // add funds to specific address
+        let subaddress_index = 3;
+        let subaddress = account_key.subaddress(subaddress_index);
+        let b58_pub_address =
+            b58_encode_public_address(&subaddress).expect("Could not encode public address");
+        service
+            .import_subaddresses(
+                &account_id.to_string(),
+                [(
+                    b58_pub_address.clone(),
+                    subaddress_index,
+                    "cheese".to_string(),
+                    subaddress.spend_public_key().to_owned(),
+                )]
+                .to_vec(),
+            )
+            .unwrap();
+
+        let value = 100 * MOB;
+        let tx_private_key = RistrettoPrivate::from_random(&mut rng);
+        let hint = EncryptedFogHint::fake_onetime_hint(&mut rng);
+        let fake_tx_out =
+            TxOut::new(value as u64, &main_public_address, &tx_private_key, hint).unwrap();
+        ViewOnlyTxo::create(
+            fake_tx_out.clone(),
+            value,
+            Some(subaddress_index),
+            Some(current_block_height),
+            &account_id.to_string(),
+            &conn,
+        )
+        .unwrap();
+
+        let balance: Balance = service
+            .get_balance_for_view_only_address(&b58_pub_address)
+            .unwrap();
+        assert_eq!(balance.unspent as u64, 100 * MOB);
+        // view only accounts have no spendable MOB
+        assert_eq!(balance.max_spendable, 0);
+        assert_eq!(balance.spent, 0);
+        assert_eq!(balance.pending, 0);
+        assert_eq!(balance.secreted, 0);
+        assert_eq!(balance.orphaned, 0);
+    }
 }