Skip to content

Commit

Permalink
zcash_client_sqlite: Generalise wallet::scanning tests
Browse files Browse the repository at this point in the history
  • Loading branch information
str4d committed Mar 9, 2024
1 parent 399ae76 commit 9a0f6cc
Show file tree
Hide file tree
Showing 3 changed files with 110 additions and 74 deletions.
6 changes: 5 additions & 1 deletion zcash_client_sqlite/src/testing/pool.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
use std::{convert::Infallible, num::NonZeroU32};

use incrementalmerkletree::Level;
use rusqlite::params;
use secrecy::Secret;
use shardtree::error::ShardTreeError;
Expand All @@ -29,7 +30,7 @@ use zcash_client_backend::{
chain::CommitmentTreeRoot,
error::Error,
wallet::input_selection::{GreedyInputSelector, GreedyInputSelectorError},
AccountBirthday, DecryptedTransaction, Ratio, WalletRead, WalletWrite,
AccountBirthday, DecryptedTransaction, Ratio, WalletRead, WalletSummary, WalletWrite,
},
decrypt_transaction,
fees::{fixed, standard, DustOutputPolicy},
Expand Down Expand Up @@ -87,13 +88,16 @@ pub(crate) trait ShieldedPoolTester {
fn fvks_equal(a: &Self::Fvk, b: &Self::Fvk) -> bool;

fn empty_tree_leaf() -> Self::MerkleTreeHash;
fn empty_tree_root(level: Level) -> Self::MerkleTreeHash;

fn put_subtree_roots<Cache>(
st: &mut TestState<Cache>,
start_index: u64,
roots: &[CommitmentTreeRoot<Self::MerkleTreeHash>],
) -> Result<(), ShardTreeError<commitment_tree::Error>>;

fn next_subtree_index(s: &WalletSummary<AccountId>) -> u64;

fn select_spendable_notes<Cache>(
st: &TestState<Cache>,
account: AccountId,
Expand Down
14 changes: 12 additions & 2 deletions zcash_client_sqlite/src/wallet/sapling.rs
Original file line number Diff line number Diff line change
Expand Up @@ -464,7 +464,7 @@ pub(crate) fn put_received_note<T: ReceivedSaplingOutput>(

#[cfg(test)]
pub(crate) mod tests {
use incrementalmerkletree::Hashable;
use incrementalmerkletree::{Hashable, Level};
use shardtree::error::ShardTreeError;
use zcash_proofs::prover::LocalTxProver;

Expand All @@ -486,7 +486,9 @@ pub(crate) mod tests {

use zcash_client_backend::{
address::Address,
data_api::{chain::CommitmentTreeRoot, DecryptedTransaction, WalletCommitmentTrees},
data_api::{
chain::CommitmentTreeRoot, DecryptedTransaction, WalletCommitmentTrees, WalletSummary,
},
keys::UnifiedSpendingKey,
wallet::{Note, ReceivedNote},
ShieldedProtocol,
Expand Down Expand Up @@ -544,6 +546,10 @@ pub(crate) mod tests {
sapling::Node::empty_leaf()
}

fn empty_tree_root(level: Level) -> Self::MerkleTreeHash {
sapling::Node::empty_root(level)
}

fn put_subtree_roots<Cache>(
st: &mut TestState<Cache>,
start_index: u64,
Expand All @@ -553,6 +559,10 @@ pub(crate) mod tests {
.put_sapling_subtree_roots(start_index, roots)
}

fn next_subtree_index(s: &WalletSummary<AccountId>) -> u64 {
s.next_sapling_subtree_index()
}

fn select_spendable_notes<Cache>(
st: &TestState<Cache>,
account: AccountId,
Expand Down
164 changes: 93 additions & 71 deletions zcash_client_sqlite/src/wallet/scanning.rs
Original file line number Diff line number Diff line change
Expand Up @@ -567,34 +567,39 @@ pub(crate) mod tests {
};

#[test]
fn scan_complete() {
fn sapling_scan_complete() {
scan_complete::<SaplingPoolTester>();
}

fn scan_complete<T: ShieldedPoolTester>() {
use ScanPriority::*;

let mut st = TestBuilder::new()
.with_block_cache()
.with_test_account(AccountBirthday::from_sapling_activation)
.build();

let dfvk = st.test_account_sapling().unwrap();
let dfvk = T::test_account_fvk(&st);
let sapling_activation_height = st.sapling_activation_height();

assert_matches!(
// In the following, we don't care what the root hashes are, they just need to be
// distinct.
st.wallet_mut().put_sapling_subtree_roots(
T::put_subtree_roots(
&mut st,
0,
&[
CommitmentTreeRoot::from_parts(
sapling_activation_height + 100,
Node::empty_root(Level::from(0))
T::empty_tree_root(Level::from(0))
),
CommitmentTreeRoot::from_parts(
sapling_activation_height + 200,
Node::empty_root(Level::from(1))
T::empty_tree_root(Level::from(1))
),
CommitmentTreeRoot::from_parts(
sapling_activation_height + 300,
Node::empty_root(Level::from(2))
T::empty_tree_root(Level::from(2))
),
]
),
Expand Down Expand Up @@ -687,7 +692,7 @@ pub(crate) mod tests {
st.wallet()
.get_wallet_summary(0)
.unwrap()
.map(|s| s.next_sapling_subtree_index()),
.map(|s| T::next_subtree_index(&s)),
Some(2),
);
}
Expand Down Expand Up @@ -725,10 +730,14 @@ pub(crate) mod tests {
}

#[test]
fn create_account_creates_ignored_range() {
fn sapling_create_account_creates_ignored_range() {
create_account_creates_ignored_range::<SaplingPoolTester>();
}

fn create_account_creates_ignored_range<T: ShieldedPoolTester>() {
use ScanPriority::*;

let (st, _, birthday, sap_active) = test_with_canopy_birthday::<SaplingPoolTester>();
let (st, _, birthday, sap_active) = test_with_canopy_birthday::<T>();
let birthday_height = birthday.height().into();

let expected = vec![
Expand Down Expand Up @@ -784,10 +793,14 @@ pub(crate) mod tests {
}

#[test]
fn update_chain_tip_with_no_subtree_roots() {
fn sapling_update_chain_tip_with_no_subtree_roots() {
update_chain_tip_with_no_subtree_roots::<SaplingPoolTester>();
}

fn update_chain_tip_with_no_subtree_roots<T: ShieldedPoolTester>() {
use ScanPriority::*;

let (mut st, _, birthday, sap_active) = test_with_canopy_birthday::<SaplingPoolTester>();
let (mut st, _, birthday, sap_active) = test_with_canopy_birthday::<T>();

// Set up the following situation:
//
Expand Down Expand Up @@ -815,10 +828,14 @@ pub(crate) mod tests {
}

#[test]
fn update_chain_tip_when_never_scanned() {
fn sapling_update_chain_tip_when_never_scanned() {
update_chain_tip_when_never_scanned::<SaplingPoolTester>();
}

fn update_chain_tip_when_never_scanned<T: ShieldedPoolTester>() {
use ScanPriority::*;

let (mut st, _, birthday, sap_active) = test_with_canopy_birthday::<SaplingPoolTester>();
let (mut st, _, birthday, sap_active) = test_with_canopy_birthday::<T>();

// Set up the following situation:
//
Expand All @@ -829,16 +846,16 @@ pub(crate) mod tests {

// Set up some shard root history before the wallet birthday.
let last_shard_start = birthday.height() - 1000;
st.wallet_mut()
.put_sapling_subtree_roots(
0,
&[CommitmentTreeRoot::from_parts(
last_shard_start,
// fake a hash, the value doesn't matter
Node::empty_leaf(),
)],
)
.unwrap();
T::put_subtree_roots(
&mut st,
0,
&[CommitmentTreeRoot::from_parts(
last_shard_start,
// fake a hash, the value doesn't matter
T::empty_tree_leaf(),
)],
)
.unwrap();

// Update the chain tip.
let tip_height = prior_tip_height + 500;
Expand All @@ -860,11 +877,15 @@ pub(crate) mod tests {
}

#[test]
fn update_chain_tip_unstable_max_scanned() {
fn sapling_update_chain_tip_unstable_max_scanned() {
update_chain_tip_unstable_max_scanned::<SaplingPoolTester>();
}

fn update_chain_tip_unstable_max_scanned<T: ShieldedPoolTester>() {
use ScanPriority::*;

// this birthday is 1234 notes into the second shard
let (mut st, dfvk, birthday, sap_active) = test_with_canopy_birthday::<SaplingPoolTester>();
let (mut st, dfvk, birthday, sap_active) = test_with_canopy_birthday::<T>();

// Set up the following situation:
//
Expand All @@ -876,16 +897,16 @@ pub(crate) mod tests {

// Set up some shard root history before the wallet birthday.
let initial_shard_end = birthday.height() - 1000;
st.wallet_mut()
.put_sapling_subtree_roots(
0,
&[CommitmentTreeRoot::from_parts(
initial_shard_end,
// fake a hash, the value doesn't matter
Node::empty_leaf(),
)],
)
.unwrap();
T::put_subtree_roots(
&mut st,
0,
&[CommitmentTreeRoot::from_parts(
initial_shard_end,
// fake a hash, the value doesn't matter
T::empty_tree_leaf(),
)],
)
.unwrap();

// Set up prior chain state. This simulates us having imported a wallet
// with a birthday 520 blocks below the chain tip.
Expand Down Expand Up @@ -928,16 +949,16 @@ pub(crate) mod tests {
// Now simulate shutting down, and then restarting 90 blocks later, after a shard
// has been completed.
let last_shard_start = prior_tip + 70;
st.wallet_mut()
.put_sapling_subtree_roots(
0,
&[CommitmentTreeRoot::from_parts(
last_shard_start,
// fake a hash, the value doesn't matter
Node::empty_leaf(),
)],
)
.unwrap();
T::put_subtree_roots(
&mut st,
0,
&[CommitmentTreeRoot::from_parts(
last_shard_start,
// fake a hash, the value doesn't matter
T::empty_tree_leaf(),
)],
)
.unwrap();

let new_tip = last_shard_start + 20;
st.wallet_mut().update_chain_tip(new_tip).unwrap();
Expand Down Expand Up @@ -972,10 +993,14 @@ pub(crate) mod tests {
}

#[test]
fn update_chain_tip_stable_max_scanned() {
fn sapling_update_chain_tip_stable_max_scanned() {
update_chain_tip_stable_max_scanned::<SaplingPoolTester>();
}

fn update_chain_tip_stable_max_scanned<T: ShieldedPoolTester>() {
use ScanPriority::*;

let (mut st, dfvk, birthday, sap_active) = test_with_canopy_birthday::<SaplingPoolTester>();
let (mut st, dfvk, birthday, sap_active) = test_with_canopy_birthday::<T>();

// Set up the following situation:
//
Expand All @@ -988,16 +1013,16 @@ pub(crate) mod tests {

// Set up some shard root history before the wallet birthday.
let second_to_last_shard_start = birthday.height() - 1000;
st.wallet_mut()
.put_sapling_subtree_roots(
0,
&[CommitmentTreeRoot::from_parts(
second_to_last_shard_start,
// fake a hash, the value doesn't matter
Node::empty_leaf(),
)],
)
.unwrap();
T::put_subtree_roots(
&mut st,
0,
&[CommitmentTreeRoot::from_parts(
second_to_last_shard_start,
// fake a hash, the value doesn't matter
T::empty_tree_leaf(),
)],
)
.unwrap();

// We have scan ranges and a subtree, but have scanned no blocks.
let summary = st.get_wallet_summary(1);
Expand Down Expand Up @@ -1032,10 +1057,7 @@ pub(crate) mod tests {
// We have scanned a block, so we now have a starting tree position, 500 blocks above the
// wallet birthday but before the end of the shard.
let summary = st.get_wallet_summary(1);
assert_eq!(
summary.as_ref().map(|s| s.next_sapling_subtree_index()),
Some(0),
);
assert_eq!(summary.as_ref().map(|s| T::next_subtree_index(s)), Some(0),);
assert_eq!(
summary.and_then(|s| s.scan_progress()),
Some(Ratio::new(1, 0x1 << SAPLING_SHARD_HEIGHT))
Expand All @@ -1044,16 +1066,16 @@ pub(crate) mod tests {
// Now simulate shutting down, and then restarting 70 blocks later, after a shard
// has been completed.
let last_shard_start = prior_tip + 50;
st.wallet_mut()
.put_sapling_subtree_roots(
0,
&[CommitmentTreeRoot::from_parts(
last_shard_start,
// fake a hash, the value doesn't matter
Node::empty_leaf(),
)],
)
.unwrap();
T::put_subtree_roots(
&mut st,
0,
&[CommitmentTreeRoot::from_parts(
last_shard_start,
// fake a hash, the value doesn't matter
T::empty_tree_leaf(),
)],
)
.unwrap();

let new_tip = last_shard_start + 20;
st.wallet_mut().update_chain_tip(new_tip).unwrap();
Expand Down

0 comments on commit 9a0f6cc

Please sign in to comment.