zcash_client_backend/data_api/testing/pool/
dsl.rs

1//! A convenient DSL for writing wallet tests.
2
3use std::{
4    convert::Infallible,
5    marker::PhantomData,
6    ops::{Deref, DerefMut},
7};
8
9use assert_matches::assert_matches;
10use nonempty::NonEmpty;
11use zcash_keys::address::Address;
12use zcash_primitives::{block::BlockHash, transaction::fees::zip317};
13use zcash_protocol::{
14    PoolType, TxId, consensus::BlockHeight, local_consensus::LocalNetwork, value::Zatoshis,
15};
16
17use zip321::Payment;
18
19use crate::{
20    data_api::{
21        Account, AccountBalance, InputSource, WalletRead, WalletTest,
22        chain::ScanSummary,
23        testing::{
24            AddressType, DataStoreFactory, FakeCompactOutput, TestAccount, TestBuilder, TestCache,
25            TestFvk, TestState, single_output_change_strategy,
26        },
27        wallet::{
28            ConfirmationsPolicy, LockRequest,
29            input_selection::{GreedyInputSelector, SpendPolicy},
30            propose_transfer,
31        },
32    },
33    fees::StandardFeeRule,
34    proposal::Proposal,
35    wallet::{LockOwner, OutputRef, OvkPolicy},
36};
37
38use super::ShieldedPoolTester;
39
40/// A type-state wrapper struct that provides convenience methods.
41pub struct TestDsl<T> {
42    /// Either a `TestBuilder<Cache, DataStoreFactory>` or
43    /// `TestState<Cache, DataStore, LocalNetwork>`.
44    inner: T,
45}
46
47impl<T> Deref for TestDsl<T> {
48    type Target = T;
49
50    fn deref(&self) -> &Self::Target {
51        &self.inner
52    }
53}
54
55impl<T> DerefMut for TestDsl<T> {
56    fn deref_mut(&mut self) -> &mut Self::Target {
57        &mut self.inner
58    }
59}
60
61impl<T> From<T> for TestDsl<T> {
62    fn from(inner: T) -> Self {
63        Self { inner }
64    }
65}
66
67impl<T> TestDsl<T> {
68    /// Perform a state transition that may change the inner type.
69    pub fn map<X>(self, f: impl FnOnce(T) -> X) -> TestDsl<X> {
70        f(self.inner).into()
71    }
72}
73
74/// [`TestDsl`] provides convenience methods for common [`TestBuilder`] scenarios.
75impl<Cache, Dsf> TestDsl<TestBuilder<Cache, Dsf>>
76where
77    Dsf: DataStoreFactory,
78{
79    /// Equip the inner [`TestBuilder`] with the provided [`DataStoreFactory`]
80    /// and [`TestCache`], as well as an account that has a birthday at Sapling
81    /// activation.
82    ///
83    /// Shorthand for the following:
84    /// ```rust,ignore
85    /// let dsl: TestDsl<TestBuilder<_, _>> = TestBuilder::new()
86    ///     .with_data_store_factory(dsf)
87    ///     .with_block_cache(tc)
88    ///     .with_account_from_sapling_activation(BlockHash([0; 32]))
89    ///     .into();
90    /// ```
91    pub fn with_sapling_birthday_account(dsf: Dsf, tc: Cache) -> Self {
92        TestBuilder::new()
93            .with_data_store_factory(dsf)
94            .with_block_cache(tc)
95            .with_account_from_sapling_activation(BlockHash([0; 32]))
96            .into()
97    }
98
99    /// Build the builder, wrapping the resulting [`TestState`] in a [`TestDsl`] and [`TestScenario`].
100    pub fn build<T: ShieldedPoolTester>(self) -> TestDsl<TestScenario<T, Cache, Dsf>> {
101        let state = self.inner.build();
102        TestScenario {
103            state,
104            _phantom: PhantomData,
105        }
106        .into()
107    }
108}
109
110/// A proxy for `FakeCompactOutput` that allows test code to omit the `fvk` and
111/// `address_type` fields, which can be derived from the `TestState` in most cases.
112pub struct TestNoteConfig<T: ShieldedPoolTester> {
113    /// The amount of the note.
114    pub value: Zatoshis,
115    /// Diversifiable full viewing key of the recipient.
116    pub fvk: Option<T::Fvk>,
117    /// Address type of the recipient.
118    pub address_type: Option<AddressType>,
119}
120
121impl<T: ShieldedPoolTester> From<Zatoshis> for TestNoteConfig<T> {
122    fn from(value: Zatoshis) -> Self {
123        TestNoteConfig {
124            value,
125            fvk: None,
126            address_type: None,
127        }
128    }
129}
130
131impl<T: ShieldedPoolTester> TestNoteConfig<T> {
132    pub fn with_address_type(mut self, address_type: AddressType) -> Self {
133        self.address_type = Some(address_type);
134        self
135    }
136
137    pub fn with_fvk(mut self, fvk: T::Fvk) -> Self {
138        self.fvk = Some(fvk);
139        self
140    }
141}
142
143pub struct AddFundsStepResult<T: ShieldedPoolTester, C: TestCache> {
144    pub block_height: BlockHeight,
145    pub insert_result: C::InsertResult,
146    /// Empty when the step was to generate an empty block
147    pub nullifiers: Vec<<T::Fvk as TestFvk>::Nullifier>,
148}
149
150/// The input and output of one "add funds" step.
151pub struct AddFundsStep<T: ShieldedPoolTester, C: TestCache> {
152    pub notes: Vec<TestNoteConfig<T>>,
153    pub results: AddFundsStepResult<T, C>,
154}
155
156/// A collection of results from adding funds to a `TestState`.
157pub struct AddFundsSummary<T: ShieldedPoolTester, C: TestCache> {
158    pub steps: Vec<AddFundsStep<T, C>>,
159    pub scan_summary: Option<ScanSummary>,
160}
161
162impl<T: ShieldedPoolTester, C: TestCache> Default for AddFundsSummary<T, C> {
163    fn default() -> Self {
164        Self {
165            steps: Default::default(),
166            scan_summary: None,
167        }
168    }
169}
170
171impl<T: ShieldedPoolTester, C: TestCache> AddFundsSummary<T, C> {
172    /// Return the first block height.
173    pub fn first_block_height(&self) -> Option<BlockHeight> {
174        self.steps.first().map(|step| step.results.block_height)
175    }
176
177    /// Return the latest block height after generating the blocks
178    /// that added funds.
179    pub fn block_height(&self) -> Option<BlockHeight> {
180        self.steps.last().map(|step| step.results.block_height)
181    }
182}
183
184#[repr(transparent)]
185pub struct TestScenario<T: ShieldedPoolTester, Cache, Dsf: DataStoreFactory> {
186    /// The current scenario state.
187    state: TestState<Cache, Dsf::DataStore, LocalNetwork>,
188    _phantom: PhantomData<T>,
189}
190
191impl<T, C, D> Deref for TestScenario<T, C, D>
192where
193    T: ShieldedPoolTester,
194    D: DataStoreFactory,
195{
196    type Target = TestState<C, D::DataStore, LocalNetwork>;
197
198    fn deref(&self) -> &Self::Target {
199        &self.state
200    }
201}
202
203impl<T, C, D> DerefMut for TestScenario<T, C, D>
204where
205    T: ShieldedPoolTester,
206    D: DataStoreFactory,
207{
208    fn deref_mut(&mut self) -> &mut Self::Target {
209        &mut self.state
210    }
211}
212
213/// Add funds scenarios.
214impl<Cache, Dsf, T> TestScenario<T, Cache, Dsf>
215where
216    T: ShieldedPoolTester,
217    Cache: TestCache,
218    Dsf: DataStoreFactory,
219{
220    /// Return the current test account balance, if possible.
221    ///
222    /// Returns `None` when no account summary data is available, which is
223    /// the case before the wallet has scanned any blocks.
224    pub fn get_account_balance(
225        &self,
226        confirmations_policy: ConfirmationsPolicy,
227    ) -> Option<AccountBalance> {
228        let account = self.get_account();
229        let binding = self
230            .wallet()
231            .get_wallet_summary(confirmations_policy)
232            .unwrap()?;
233        let balance = binding.account_balances().get(&account.id())?;
234        Some(*balance)
235    }
236
237    /// Adds funds from a single note from an address of the given type.
238    ///
239    /// Returns the current block height, cache insert result and test viewing key nullifier.
240    ///
241    /// This is shorthand for:
242    /// ```rust,ignore
243    /// {
244    ///     let dfvk = T::test_account_fvk(&st);
245    ///     let output@(h, _, _) = st.generate_next_block(&dfvk, address_type, zatoshis);
246    ///     st.scan_cached_blocks(h, 1);
247    ///     output
248    /// }
249    /// ```
250    ///
251    /// This also verifies that the test account contains the expected funds as
252    /// part of the _total_, and that the funds are spendable with the minimum number
253    /// of confirmations.
254    pub fn add_a_single_note_checking_balance(
255        &mut self,
256        note: impl Into<TestNoteConfig<T>>,
257    ) -> (
258        BlockHeight,
259        Cache::InsertResult,
260        <T::Fvk as TestFvk>::Nullifier,
261    ) {
262        let mut summary = self.add_notes_checking_balance([[note]]);
263        let res = summary.steps.pop().unwrap().results;
264        (res.block_height, res.insert_result, res.nullifiers[0])
265    }
266
267    fn scanned_block_height(&self) -> BlockHeight {
268        self.wallet()
269            .block_max_scanned()
270            .unwrap()
271            .map(|meta| meta.block_height())
272            .unwrap_or_else(|| BlockHeight::from_u32(0))
273    }
274
275    /// Generates `N` empty blocks and scans them.
276    ///
277    /// Returns the current block height.
278    pub fn add_empty_blocks(&mut self, n: usize) -> BlockHeight {
279        let mut out_height = self.scanned_block_height();
280        for _ in 0..n {
281            let (h, _) = self.generate_empty_block();
282            out_height = h;
283            self.scan_cached_blocks(h, 1);
284        }
285        out_height
286    }
287
288    /// Returns the test account.
289    pub fn get_account(&self) -> TestAccount<Dsf::Account> {
290        self.test_account().expect("not configured").clone()
291    }
292
293    /// Creates a `FakeCompactOutput` from the given `TestNoteConfig`.
294    fn make_fake_output(&self, note_config: &TestNoteConfig<T>) -> FakeCompactOutput<T::Fvk> {
295        let TestNoteConfig {
296            value,
297            fvk,
298            address_type,
299        } = note_config;
300        FakeCompactOutput::new(
301            fvk.clone().unwrap_or_else(|| T::test_account_fvk(self)),
302            address_type.unwrap_or(AddressType::DefaultExternal),
303            *value,
304        )
305    }
306
307    /// Add funds from multiple notes in one or more blocks, or generate empty blocks.
308    ///
309    /// This step also verifies that the test account contains the expected
310    /// funds as part of the _total_. Keep in mind that these funds may not yet
311    /// be _spendable_ due to the number of confirmations required.
312    ///
313    /// Returns a summary of steps.
314    ///
315    /// ## Parameters
316    ///
317    /// * `blocks` - A collection of "blocks", where each "block" is a collection of
318    ///   "notes". More specifically, "notes" can be anything that can be converted
319    ///   into a [`TestNoteConfig`]. This allows you to add multiple blocks that each
320    ///   containing zero or more notes with one call to `add_notes`.
321    ///
322    /// ## Note
323    /// Keep in mind:
324    /// * Each block coalesces these notes into a single transaction.
325    /// * Funds are added to the default test account.
326    pub fn add_notes_checking_balance(
327        &mut self,
328        blocks: impl IntoIterator<Item = impl IntoIterator<Item = impl Into<TestNoteConfig<T>>>>,
329    ) -> AddFundsSummary<T, Cache> {
330        let mut from_height = None;
331        let mut current_height = self.scanned_block_height();
332        let mut limit = 0;
333        let account = self.get_account();
334        let starting_balance = self
335            .get_account_balance(ConfirmationsPolicy::MIN)
336            .map(|b| b.total())
337            .unwrap_or(Zatoshis::ZERO);
338        let mut expected_total = starting_balance;
339        let mut summary = AddFundsSummary::default();
340        for notes in blocks.into_iter() {
341            // Collect the notes while also counting their value.
342            let (fake_outputs, note_configs): (Vec<_>, Vec<_>) = notes
343                .into_iter()
344                .map(|into_note_config| {
345                    let note_config = into_note_config.into();
346                    if note_config.value > zip317::MARGINAL_FEE {
347                        // Don't include uneconomic (dust) notes in the expected
348                        // total, as the balance won't include them.
349                        expected_total = (expected_total + note_config.value).unwrap();
350                    }
351                    (self.make_fake_output(&note_config), note_config)
352                })
353                .unzip();
354            let step_result = if fake_outputs.is_empty() {
355                let (h, r) = self.generate_empty_block();
356                AddFundsStepResult {
357                    block_height: h,
358                    insert_result: r,
359                    nullifiers: vec![],
360                }
361            } else {
362                let (h, r, n) = self.generate_next_block_multi(&fake_outputs);
363                AddFundsStepResult {
364                    block_height: h,
365                    insert_result: r,
366                    nullifiers: n,
367                }
368            };
369            current_height = step_result.block_height;
370            if from_height.is_none() {
371                from_height = Some(current_height);
372            }
373            summary.steps.push(AddFundsStep {
374                notes: note_configs,
375                results: step_result,
376            });
377            limit += 1;
378        }
379        if let Some(from_height) = from_height {
380            summary.scan_summary = Some(self.scan_cached_blocks(from_height, limit));
381        }
382
383        // Do most of the assertions that we care about at the "add funds" callsites
384        assert_eq!(
385            self.get_total_balance(account.id()),
386            expected_total,
387            "Unexpected total balance"
388        );
389        assert_eq!(
390            self.wallet()
391                .block_max_scanned()
392                .unwrap()
393                .unwrap()
394                .block_height(),
395            current_height
396        );
397        assert_eq!(
398            self.get_spendable_balance(account.id(), ConfirmationsPolicy::MIN),
399            expected_total
400        );
401
402        summary
403    }
404
405    /// Proposes a ZIP 317 transfer of `amount` zatoshis from the test account to
406    /// `to`, panicking if proposal construction fails.
407    ///
408    /// This fixes the values that are constant across most transfer scenarios:
409    /// the test account, [`StandardFeeRule::Zip317`], [`ConfirmationsPolicy::MIN`],
410    /// no memos, and the tester's shielded pool as the fallback change pool.
411    pub fn propose_transfer_to(
412        &mut self,
413        to: &Address,
414        amount: Zatoshis,
415    ) -> Proposal<StandardFeeRule, <Dsf::DataStore as InputSource>::NoteRef> {
416        let account_id = self.get_account().id();
417        self.propose_standard_transfer::<Infallible>(
418            account_id,
419            StandardFeeRule::Zip317,
420            ConfirmationsPolicy::MIN,
421            to,
422            amount,
423            None,
424            None,
425            T::SHIELDED_PROTOCOL,
426        )
427        .unwrap()
428    }
429
430    /// Proposes and executes a transfer of `amount` zatoshis to `to` with
431    /// [`OvkPolicy::Sender`], returning the single resulting transaction id.
432    ///
433    /// Shorthand for [`Self::propose_transfer_to`] followed by
434    /// [`TestState::create_proposed_transactions`]. (Named `spend_to` to avoid
435    /// clashing with the lower-level [`TestState::spend`].)
436    pub fn spend_to(&mut self, to: &Address, amount: Zatoshis) -> TxId {
437        let account = self.get_account();
438        let proposal = self.propose_transfer_to(to, amount);
439        self.create_proposed_transactions::<Infallible, _, Infallible, _>(
440            account.usk(),
441            OvkPolicy::Sender,
442            &proposal,
443        )
444        .unwrap()[0]
445    }
446
447    /// Creates the transactions for `proposal` with [`OvkPolicy::Sender`],
448    /// asserts that exactly `expected` transactions were produced, and returns
449    /// their ids.
450    ///
451    /// This collapses the common
452    /// [`TestState::create_proposed_transactions`]-then-assert-count
453    /// boilerplate. It fixes the values that are constant across nearly every
454    /// execution site: the test account's spending key and
455    /// [`OvkPolicy::Sender`]. On failure it reports the full result (including
456    /// any error), matching the diagnostics of the hand-written form.
457    pub fn create_proposed_expecting(
458        &mut self,
459        proposal: &Proposal<StandardFeeRule, <Dsf::DataStore as InputSource>::NoteRef>,
460        expected: usize,
461    ) -> NonEmpty<TxId> {
462        let account = self.get_account();
463        let result = self.create_proposed_transactions::<Infallible, _, Infallible, _>(
464            account.usk(),
465            OvkPolicy::Sender,
466            proposal,
467        );
468        assert_matches!(&result, Ok(txids) if txids.len() == expected);
469        result.unwrap()
470    }
471
472    /// Asserts that proposing a transfer of `amount` zatoshis to `to` fails with
473    /// [`Error::InsufficientFunds`], reporting `available` available and exactly
474    /// `required` required.
475    ///
476    /// [`Error::InsufficientFunds`]: crate::data_api::error::Error::InsufficientFunds
477    pub fn expect_insufficient_funds(
478        &mut self,
479        to: &Address,
480        amount: Zatoshis,
481        available: Zatoshis,
482        required: Zatoshis,
483    ) {
484        self.expect_insufficient_funds_with(
485            to,
486            amount,
487            ConfirmationsPolicy::MIN,
488            available,
489            required,
490        )
491    }
492
493    /// Like [`Self::expect_insufficient_funds`], but with an explicit
494    /// `confirmations_policy` instead of the fixed [`ConfirmationsPolicy::MIN`].
495    pub fn expect_insufficient_funds_with(
496        &mut self,
497        to: &Address,
498        amount: Zatoshis,
499        confirmations_policy: ConfirmationsPolicy,
500        available: Zatoshis,
501        required: Zatoshis,
502    ) {
503        let account_id = self.get_account().id();
504        assert_matches!(
505            self.propose_standard_transfer::<Infallible>(
506                account_id,
507                StandardFeeRule::Zip317,
508                confirmations_policy,
509                to,
510                amount,
511                None,
512                None,
513                T::SHIELDED_PROTOCOL,
514            ),
515            Err(crate::data_api::error::Error::InsufficientFunds { available: a, required: r })
516                if a == available && r == required,
517            "expected InsufficientFunds (available={}, required={}) proposing {}",
518            u64::from(available),
519            u64::from(required),
520            u64::from(amount)
521        );
522    }
523
524    /// Mines a single "decoy" block that pays `value` to a throwaway external
525    /// address derived from `seed` (so the funds do not accrue to the test
526    /// account), returning the new block height. Does not scan.
527    pub fn mine_decoy_block(&mut self, seed: u8, value: Zatoshis) -> BlockHeight {
528        let (h, _, _) = self.generate_next_block(
529            &T::sk_to_fvk(&T::sk(&[seed; 32])),
530            AddressType::DefaultExternal,
531            value,
532        );
533        h
534    }
535
536    /// Mines one decoy block per `seed` (see [`Self::mine_decoy_block`]),
537    /// returning the height of the last block mined, if any. Does not scan.
538    pub fn mine_decoy_blocks(
539        &mut self,
540        seeds: impl IntoIterator<Item = u8>,
541        value: Zatoshis,
542    ) -> Option<BlockHeight> {
543        let mut last = None;
544        for seed in seeds {
545            last = Some(self.mine_decoy_block(seed, value));
546        }
547        last
548    }
549
550    /// Returns an [`OutputRef`] for the wallet's single received note in the
551    /// tester's shielded pool, asserting that exactly one such note exists.
552    pub fn sole_note_ref(&self) -> OutputRef {
553        let notes = self.wallet().get_notes(T::SHIELDED_PROTOCOL).unwrap();
554        assert_eq!(notes.len(), 1);
555        let note = &notes[0];
556        OutputRef::new(
557            *note.txid(),
558            PoolType::Shielded(note.note().pool()),
559            u32::from(note.output_index()),
560        )
561    }
562
563    /// Returns an [`OutputRef`] for the single received note whose value equals
564    /// `value`, panicking if no such note exists.
565    ///
566    /// Used by the multi-note scenarios, which fund the account with notes of
567    /// distinct values so that each note can be identified by its value.
568    pub fn note_ref_by_value(&self, value: Zatoshis) -> OutputRef {
569        let notes = self.wallet().get_notes(T::SHIELDED_PROTOCOL).unwrap();
570        let note = notes
571            .iter()
572            .find(|n| n.note().value() == value)
573            .expect("a note with the requested value exists");
574        OutputRef::new(
575            *note.txid(),
576            PoolType::Shielded(note.note().pool()),
577            u32::from(note.output_index()),
578        )
579    }
580
581    /// Proposes a ZIP 317 transfer of `amount` zatoshis to `to` that locks its
582    /// own selected inputs on behalf of `owner` for `lock_for_blocks` blocks,
583    /// panicking if proposal construction fails.
584    ///
585    /// This is the locking counterpart of [`Self::propose_transfer_to`]: it
586    /// goes through the lower-level [`propose_transfer`] so that a
587    /// [`LockRequest`] can be attached, and fixes the same constants (the test
588    /// account, a [`GreedyInputSelector`], the tester's single-output change
589    /// strategy, [`ConfirmationsPolicy::MIN`], the default [`SpendPolicy`], and
590    /// no requested transaction version).
591    pub fn propose_locking_transfer(
592        &mut self,
593        to: &Address,
594        amount: Zatoshis,
595        owner: LockOwner,
596        lock_for_blocks: u32,
597    ) -> Proposal<StandardFeeRule, <Dsf::DataStore as InputSource>::NoteRef> {
598        let account_id = self.get_account().id();
599        let input_selector = GreedyInputSelector::new();
600        let change_strategy =
601            single_output_change_strategy(StandardFeeRule::Zip317, None, T::SHIELDED_PROTOCOL);
602        let request = zip321::TransactionRequest::new(vec![Payment::without_memo(
603            to.to_zcash_address(self.network()),
604            amount,
605        )])
606        .unwrap();
607        let network = *self.network();
608        propose_transfer::<_, _, _, _, Infallible>(
609            self.wallet_mut(),
610            &network,
611            account_id,
612            &input_selector,
613            &change_strategy,
614            request,
615            ConfirmationsPolicy::MIN,
616            &SpendPolicy::default(),
617            Some(LockRequest::new(owner, lock_for_blocks)),
618            None,
619        )
620        .unwrap()
621    }
622}
zcash_client_backend/data_api/testing/pool/dsl.rs

zcash_client_backend/data_api/testing/pool/
dsl.rs
