Skip to content

Commit

Permalink
Refactor: Extract Method check_balance
Browse files Browse the repository at this point in the history
  • Loading branch information
@pacu
pacu committed Oct 26, 2024
1 parent db04621 commit ce49ed4
Showing 1 changed file with 73 additions and 115 deletions.
188 changes: 73 additions & 115 deletions zcash_client_backend/src/data_api/testing/transparent.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,9 +13,61 @@ use crate::{
use assert_matches::assert_matches;
use sapling::zip32::ExtendedSpendingKey;
use zcash_primitives::{
block::BlockHash,
transaction::components::{amount::NonNegativeAmount, OutPoint, TxOut},
block::BlockHash, legacy::TransparentAddress, transaction::components::{amount::NonNegativeAmount, OutPoint, TxOut}
};
use zcash_protocol::local_consensus::LocalNetwork;

use super::TestAccount;

fn check_balance<DSF>(st: &TestState<impl TestCache, <DSF as DataStoreFactory>::DataStore, LocalNetwork>, account: &TestAccount<<DSF as DataStoreFactory>::Account>, taddr: &TransparentAddress, min_confirmations: u32, expected: &Balance)
where
DSF: DataStoreFactory,
{
// Check the wallet summary returns the expected transparent balance.
let summary = st
.wallet()
.get_wallet_summary(min_confirmations)
.unwrap()
.unwrap();
let balance = summary.account_balances().get(&account.id()).unwrap();

#[allow(deprecated)]
let old_unshielded_value = balance.unshielded();
assert_eq!(old_unshielded_value, expected.total());
assert_eq!(balance.unshielded_balance(), expected);

// Check the older APIs for consistency.
let mempool_height = st.wallet().chain_height().unwrap().unwrap() + 1;
assert_eq!(
st.wallet()
.get_transparent_balances(account.id(), mempool_height)
.unwrap()
.get(taddr)
.cloned()
.unwrap_or(NonNegativeAmount::ZERO),
expected.total(),
);
assert_eq!(
st.wallet()
.get_spendable_transparent_outputs(taddr, mempool_height, 0)
.unwrap()
.into_iter()
.map(|utxo| utxo.value())
.sum::<Option<NonNegativeAmount>>(),
Some(expected.spendable_value()),
);

// we currently treat min_confirmation the same regardless they are 0 or 1
assert_eq!(
st.wallet()
.get_spendable_transparent_outputs(taddr, mempool_height, 1)
.unwrap()
.into_iter()
.map(|utxo| utxo.value())
.sum::<Option<NonNegativeAmount>>(),
Some(expected.spendable_value()),
);
}

pub fn put_received_transparent_utxo<DSF>(dsf: DSF)
where
Expand Down Expand Up @@ -137,59 +189,11 @@ where
}
st.scan_cached_blocks(start_height, 10);

let check_balance = |st: &TestState<_, DSF::DataStore, _>,
min_confirmations: u32,
expected: &crate::data_api::Balance| {
// Check the wallet summary returns the expected transparent balance.
let summary = st
.wallet()
.get_wallet_summary(min_confirmations)
.unwrap()
.unwrap();
let balance = summary.account_balances().get(&account.id()).unwrap();

#[allow(deprecated)]
let old_unshielded_value = balance.unshielded();
assert_eq!(old_unshielded_value, expected.total());
assert_eq!(balance.unshielded_balance(), expected);

// Check the older APIs for consistency.
let mempool_height = st.wallet().chain_height().unwrap().unwrap() + 1;
assert_eq!(
st.wallet()
.get_transparent_balances(account.id(), mempool_height)
.unwrap()
.get(taddr)
.cloned()
.unwrap_or(NonNegativeAmount::ZERO),
expected.total(),
);
assert_eq!(
st.wallet()
.get_spendable_transparent_outputs(taddr, mempool_height, 0)
.unwrap()
.into_iter()
.map(|utxo| utxo.value())
.sum::<Option<NonNegativeAmount>>(),
Some(expected.spendable_value()),
);

// we currently treat min_confirmation the same regardless they are 0 or 1
assert_eq!(
st.wallet()
.get_spendable_transparent_outputs(taddr, mempool_height, 1)
.unwrap()
.into_iter()
.map(|utxo| utxo.value())
.sum::<Option<NonNegativeAmount>>(),
Some(expected.spendable_value()),
);
};

// The wallet starts out with zero balance.
check_balance(&st, 0, &Balance::ZERO);
check_balance::<DSF>(&st, &account, taddr,0, &Balance::ZERO);
// we currently treat min_confirmation the same regardless they are 0 or 1
check_balance(&st, 1, &Balance::ZERO);
check_balance::<DSF>(&st, &account, taddr,1, &Balance::ZERO);

// Create a fake transparent output.
let value = NonNegativeAmount::from_u64(100000).unwrap();
Expand All @@ -211,8 +215,9 @@ where
// add the spendable value to the expected balance
zero_or_one_conf_value.add_spendable_value(value).unwrap();

check_balance(&st, 0, &zero_or_one_conf_value);
check_balance(&st, 1, &zero_or_one_conf_value);
check_balance::<DSF>(&st, &account, taddr,0, &zero_or_one_conf_value);
// we currently treat min_confirmation the same regardless they are 0 or 1
check_balance::<DSF>(&st, &account, taddr,1, &zero_or_one_conf_value);

// Shield the output.
let input_selector = GreedyInputSelector::new();
Expand All @@ -236,18 +241,18 @@ where

// The wallet should have zero transparent balance, because the shielding
// transaction can be mined.
check_balance(&st, 0, &Balance::ZERO);
check_balance::<DSF>(&st, &account, taddr,0, &Balance::ZERO);
// we currently treat min_confirmation the same regardless they are 0 or 1
check_balance(&st, 1, &Balance::ZERO);
check_balance::<DSF>(&st, &account, taddr,1, &Balance::ZERO);

// Mine the shielding transaction.
let (mined_height, _) = st.generate_next_block_including(txid);
st.scan_cached_blocks(mined_height, 1);

// The wallet should still have zero transparent balance.
check_balance(&st, 0, &Balance::ZERO);
check_balance::<DSF>(&st, &account, taddr,0, &Balance::ZERO);
// we currently treat min_confirmation the same regardless they are 0 or 1
check_balance(&st, 1, &Balance::ZERO);
check_balance::<DSF>(&st, &account, taddr,1, &Balance::ZERO);

// Unmine the shielding transaction via a reorg.
st.wallet_mut()
Expand All @@ -256,9 +261,9 @@ where
assert_eq!(st.wallet().chain_height().unwrap(), Some(mined_height - 1));

// The wallet should still have zero transparent balance.
check_balance(&st, 0, &Balance::ZERO);
check_balance::<DSF>(&st, &account, taddr,0, &Balance::ZERO);
// we currently treat min_confirmation the same regardless they are 0 or 1
check_balance(&st, 1, &Balance::ZERO);
check_balance::<DSF>(&st, &account, taddr,1, &Balance::ZERO);

// Expire the shielding transaction.
let expiry_height = st
Expand All @@ -269,9 +274,9 @@ where
.expiry_height();
st.wallet_mut().update_chain_tip(expiry_height).unwrap();

check_balance(&st, 0, &zero_or_one_conf_value);
check_balance::<DSF>(&st, &account, taddr, 0, &zero_or_one_conf_value);
// we currently treat min_confirmation the same regardless they are 0 or 1
check_balance(&st, 1, &zero_or_one_conf_value);
check_balance::<DSF>(&st, &account, taddr,1, &zero_or_one_conf_value);
}

/// This test attempts to verify that transparent funds spendability is
Expand Down Expand Up @@ -304,57 +309,10 @@ where
}
st.scan_cached_blocks(start_height, 10);

// TODO: abstract this function so it's not duplicated among different tests.
let check_balance = |st: &TestState<_, DSF::DataStore, _>,
min_confirmations: u32,
expected: &crate::data_api::Balance| {
// Check the wallet summary returns the expected transparent balance.
let summary = st
.wallet()
.get_wallet_summary(min_confirmations)
.unwrap()
.unwrap();
let balance = summary.account_balances().get(&account.id()).unwrap();

assert_eq!(balance.unshielded_balance(), expected);

// Check the older APIs for consistency.
let mempool_height = st.wallet().chain_height().unwrap().unwrap() + 1;
assert_eq!(
st.wallet()
.get_transparent_balances(account.id(), mempool_height)
.unwrap()
.get(taddr)
.cloned()
.unwrap_or(NonNegativeAmount::ZERO),
expected.total(),
);
assert_eq!(
st.wallet()
.get_spendable_transparent_outputs(taddr, mempool_height, 0)
.unwrap()
.into_iter()
.map(|utxo| utxo.value())
.sum::<Option<NonNegativeAmount>>(),
Some(expected.spendable_value()),
);

// we currently treat min_confirmation the same regardless they are 0 or 1
assert_eq!(
st.wallet()
.get_spendable_transparent_outputs(taddr, mempool_height, 1)
.unwrap()
.into_iter()
.map(|utxo| utxo.value())
.sum::<Option<NonNegativeAmount>>(),
Some(expected.spendable_value()),
);
};

// The wallet starts out with zero balance.
check_balance(&st, 0, &Balance::ZERO);
check_balance::<DSF>(&st as &TestState<_, DSF::DataStore, _>, &account, taddr, 0, &Balance::ZERO);
// we currently treat min_confirmation the same regardless they are 0 or 1
check_balance(&st, 1, &Balance::ZERO);
check_balance::<DSF>(&st, &account, taddr, 1, &Balance::ZERO);

// Create a fake transparent output.
let value = NonNegativeAmount::from_u64(100000).unwrap();
Expand All @@ -376,21 +334,21 @@ where
// add the spendable value to the expected balance
zero_or_one_conf_value.add_spendable_value(value).unwrap();

check_balance(&st, 0, &zero_or_one_conf_value);
check_balance::<DSF>(&st, &account, taddr,0, &zero_or_one_conf_value);
// we currently treat min_confirmation the same regardless they are 0 or 1
check_balance(&st, 1, &zero_or_one_conf_value);
check_balance::<DSF>(&st, &account, taddr,1, &zero_or_one_conf_value);

// now if we increase the number of confirmations our spendable balance should
// be zero and the total balance equal to `value`
let mut not_confirmed_yet_value = Balance::ZERO;

not_confirmed_yet_value.add_pending_spendable_value(value).unwrap();

check_balance(&st, 2, &not_confirmed_yet_value);
check_balance::<DSF>(&st, &account, taddr,2, &not_confirmed_yet_value);

st.generate_empty_block();

// now we generate one more block and the balance should be the same as when the
// check_balance function was called with zero or one confirmation.
check_balance(&st, 2, &zero_or_one_conf_value);
check_balance::<DSF>(&st, &account, taddr,2, &zero_or_one_conf_value);
}

0 comments on commit ce49ed4

Please sign in to comment.