zebra_state/service/finalized_state/zebra_db/
arbitrary.rs

1//! Arbitrary value generation and test harnesses for high-level typed database access.
2
3#![allow(unused_imports)]
4
5use std::ops::Deref;
6
7use zebra_chain::{amount::NonNegative, block::Block, sprout, value_balance::ValueBalance};
8
9use crate::service::finalized_state::{
10    disk_db::{DiskDb, DiskWriteBatch, WriteDisk},
11    ZebraDb,
12};
13
14// Enable older test code to automatically access the inner database via Deref coercion.
15impl Deref for ZebraDb {
16    type Target = DiskDb;
17
18    fn deref(&self) -> &Self::Target {
19        self.db()
20    }
21}
22
23impl ZebraDb {
24    /// Returns the inner database.
25    ///
26    /// This is a test-only method, because it allows write access and raw read
27    /// access to the RocksDB instance.
28    pub fn db(&self) -> &DiskDb {
29        &self.db
30    }
31
32    /// Allow to set up a fake value pool in the database for testing purposes.
33    #[cfg(any(test, feature = "proptest-impl"))]
34    pub fn set_finalized_value_pool(&self, fake_value_pool: ValueBalance<NonNegative>) {
35        let mut batch = DiskWriteBatch::new();
36        let value_pool_cf = self.db().cf_handle("tip_chain_value_pool").unwrap();
37
38        batch.zs_insert(&value_pool_cf, (), fake_value_pool);
39        self.db().write(batch).unwrap();
40    }
41
42    /// Artificially prime the note commitment tree anchor sets with anchors
43    /// referenced in a block, for testing purposes _only_.
44    #[cfg(any(test, feature = "proptest-impl"))]
45    pub fn populate_with_anchors(&self, block: &Block) {
46        let mut batch = DiskWriteBatch::new();
47
48        let sprout_anchors = self.db().cf_handle("sprout_anchors").unwrap();
49        let sapling_anchors = self.db().cf_handle("sapling_anchors").unwrap();
50        let orchard_anchors = self.db().cf_handle("orchard_anchors").unwrap();
51
52        let sprout_tree = sprout::tree::NoteCommitmentTree::default();
53        // Calculate the root so we pass the tree with a cached root to the database. We need to do
54        // that because the database checks if the tree has a cached root.
55        sprout_tree.root();
56
57        for transaction in block.transactions.iter() {
58            // Sprout
59            for joinsplit in transaction.sprout_groth16_joinsplits() {
60                batch.zs_insert(&sprout_anchors, joinsplit.anchor, sprout_tree.clone());
61            }
62
63            // Sapling
64            for anchor in transaction.sapling_anchors() {
65                batch.zs_insert(&sapling_anchors, anchor, ());
66            }
67
68            // Orchard
69            if let Some(orchard_shielded_data) = transaction.orchard_shielded_data() {
70                batch.zs_insert(&orchard_anchors, orchard_shielded_data.shared_anchor, ());
71            }
72        }
73
74        self.db().write(batch).unwrap();
75    }
76}
zebra_state/service/finalized_state/zebra_db/arbitrary.rs

zebra_state/service/finalized_state/zebra_db/
arbitrary.rs
