1//! Arbitrary value generation and test harnesses for high-level typed database access.
23#![allow(unused_imports)]
45use std::ops::Deref;
67use zebra_chain::{amount::NonNegative, block::Block, sprout, value_balance::ValueBalance};
89use crate::service::finalized_state::{
10 disk_db::{DiskDb, DiskWriteBatch, WriteDisk},
11 ZebraDb,
12};
1314// Enable older test code to automatically access the inner database via Deref coercion.
15impl Deref for ZebraDb {
16type Target = DiskDb;
1718fn deref(&self) -> &Self::Target {
19self.db()
20 }
21}
2223impl ZebraDb {
24/// Returns the inner database.
25 ///
26 /// This is a test-only and shielded-scan-only method, because it allows write access
27 /// and raw read access to the RocksDB instance.
28pub fn db(&self) -> &DiskDb {
29&self.db
30 }
3132/// Allow to set up a fake value pool in the database for testing purposes.
33#[cfg(any(test, feature = "proptest-impl"))]
34pub fn set_finalized_value_pool(&self, fake_value_pool: ValueBalance<NonNegative>) {
35let mut batch = DiskWriteBatch::new();
36let value_pool_cf = self.db().cf_handle("tip_chain_value_pool").unwrap();
3738 batch.zs_insert(&value_pool_cf, (), fake_value_pool);
39self.db().write(batch).unwrap();
40 }
4142/// 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"))]
45pub fn populate_with_anchors(&self, block: &Block) {
46let mut batch = DiskWriteBatch::new();
4748let sprout_anchors = self.db().cf_handle("sprout_anchors").unwrap();
49let sapling_anchors = self.db().cf_handle("sapling_anchors").unwrap();
50let orchard_anchors = self.db().cf_handle("orchard_anchors").unwrap();
5152let 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.
55sprout_tree.root();
5657for transaction in block.transactions.iter() {
58// Sprout
59for joinsplit in transaction.sprout_groth16_joinsplits() {
60 batch.zs_insert(&sprout_anchors, joinsplit.anchor, sprout_tree.clone());
61 }
6263// Sapling
64for anchor in transaction.sapling_anchors() {
65 batch.zs_insert(&sapling_anchors, anchor, ());
66 }
6768// Orchard
69if let Some(orchard_shielded_data) = transaction.orchard_shielded_data() {
70 batch.zs_insert(&orchard_anchors, orchard_shielded_data.shared_anchor, ());
71 }
72 }
7374self.db().write(batch).unwrap();
75 }
76}