Refactor insert_fake_orchard_shielded_data into separate V5 and V6 helpers

Author: dmidem

zebra-chain/src/orchard/shielded_data_flavor.rs

@@ -53,6 +53,7 @@ pub trait ShieldedDataFlavor: OrchardFlavor {
     /// A type representing a burn field for this protocol version.
     #[cfg(all(zcash_unstable = "nu7", feature = "tx_v6"))]
     type BurnType: Clone
+        + Default
         + Debug
         + ZcashDeserialize
         + ZcashSerialize

zebra-chain/src/transaction/arbitrary.rs

@@ -1162,22 +1162,13 @@ pub fn transactions_from_blocks<'a>(
     })
 }
 
-// FIXME: make it a generic to support V6?
-/// Modify a V5 transaction to insert fake Orchard shielded data.
-///
 /// Creates a fake instance of [`orchard::ShieldedData`] with one fake action. Note that both the
 /// action and the shielded data are invalid and shouldn't be used in tests that require them to be
 /// valid.
-///
-/// A mutable reference to the inserted shielded data is returned, so that the caller can further
-/// customize it if required.
-///
-/// # Panics
-///
-/// Panics if the transaction to be modified is not V5.
-pub fn insert_fake_orchard_shielded_data(
-    transaction: &mut Transaction,
-) -> &mut orchard::ShieldedData<orchard::OrchardVanilla> {
+pub fn create_fake_orchard_shielded_data<Flavor: orchard::ShieldedDataFlavor + 'static>(
+) -> orchard::ShieldedData<Flavor>
+//where <<Flavor as orchard::ShieldedDataFlavor>::EncryptedNote as Arbitrary>::Strategy: 'static
+{
     // Create a dummy action
     let mut runner = TestRunner::default();
     let dummy_action = orchard::Action::arbitrary()
@@ -1192,7 +1183,7 @@ pub fn insert_fake_orchard_shielded_data(
     };
 
     // Place the dummy action inside the Orchard shielded data
-    let dummy_shielded_data = orchard::ShieldedData::<orchard::OrchardVanilla> {
+    orchard::ShieldedData::<Flavor> {
         flags: orchard::Flags::empty(),
         value_balance: Amount::try_from(0).expect("invalid transaction amount"),
         shared_anchor: orchard::tree::Root::default(),
@@ -1201,15 +1192,27 @@ pub fn insert_fake_orchard_shielded_data(
         binding_sig: Signature::from([0u8; 64]),
         #[cfg(all(zcash_unstable = "nu7", feature = "tx_v6"))]
         burn: Default::default(),
-    };
+    }
+}
 
+/// Modify a V5 transaction to insert fake Orchard shielded data.
+///
+/// A mutable reference to the inserted shielded data is returned, so that the caller can further
+/// customize it if required.
+///
+/// # Panics
+///
+/// Panics if the transaction to be modified is not V5.
+pub fn insert_fake_v5_orchard_shielded_data(
+    transaction: &mut Transaction,
+) -> &mut orchard::ShieldedData<orchard::OrchardVanilla> {
     // Replace the shielded data in the transaction
     match transaction {
         Transaction::V5 {
             orchard_shielded_data,
             ..
         } => {
-            *orchard_shielded_data = Some(dummy_shielded_data);
+            *orchard_shielded_data = Some(create_fake_orchard_shielded_data());
 
             orchard_shielded_data
                 .as_mut()
@@ -1218,3 +1221,31 @@ pub fn insert_fake_orchard_shielded_data(
         _ => panic!("Fake V5 transaction is not V5"),
     }
 }
+
+/// Modify a V6 transaction to insert fake Orchard shielded data.
+///
+/// A mutable reference to the inserted shielded data is returned, so that the caller can further
+/// customize it if required.
+///
+/// # Panics
+///
+/// Panics if the transaction to be modified is not V6.
+#[cfg(all(zcash_unstable = "nu7", feature = "tx_v6"))]
+pub fn insert_fake_v6_orchard_shielded_data(
+    transaction: &mut Transaction,
+) -> &mut orchard::ShieldedData<orchard::OrchardZSA> {
+    // Replace the shielded data in the transaction
+    match transaction {
+        Transaction::V6 {
+            orchard_shielded_data,
+            ..
+        } => {
+            *orchard_shielded_data = Some(create_fake_orchard_shielded_data());
+
+            orchard_shielded_data
+                .as_mut()
+                .expect("shielded data was just inserted")
+        }
+        _ => panic!("Fake V6 transaction is not V6"),
+    }
+}

zebra-consensus/src/transaction/tests.rs

@@ -28,14 +28,17 @@ use zebra_chain::{
     sprout,
     transaction::{
         arbitrary::{
-            insert_fake_orchard_shielded_data, test_transactions, transactions_from_blocks,
+            insert_fake_v5_orchard_shielded_data, test_transactions, transactions_from_blocks,
             v5_transactions,
         },
         zip317, Hash, HashType, JoinSplitData, LockTime, Transaction,
     },
     transparent::{self, CoinbaseData, CoinbaseSpendRestriction},
 };
 
+#[cfg(all(zcash_unstable = "nu7", feature = "tx_v6"))]
+use zebra_chain::transaction::arbitrary::insert_fake_v6_orchard_shielded_data;
+
 use zebra_node_services::mempool;
 use zebra_state::ValidateContextError;
 use zebra_test::mock_service::MockService;
@@ -1218,7 +1221,7 @@ fn v5_coinbase_transaction_without_enable_spends_flag_passes_validation() {
             .find(|transaction| transaction.is_coinbase())
             .expect("V5 coinbase tx");
 
-        let shielded_data = insert_fake_orchard_shielded_data(&mut tx);
+        let shielded_data = insert_fake_v5_orchard_shielded_data(&mut tx);
 
         assert!(!shielded_data.flags.contains(Flags::ENABLE_SPENDS));
 
@@ -1233,7 +1236,7 @@ fn v5_coinbase_transaction_with_enable_spends_flag_fails_validation() {
             .find(|transaction| transaction.is_coinbase())
             .expect("V5 coinbase tx");
 
-        let shielded_data = insert_fake_orchard_shielded_data(&mut tx);
+        let shielded_data = insert_fake_v5_orchard_shielded_data(&mut tx);
 
         assert!(!shielded_data.flags.contains(Flags::ENABLE_SPENDS));
 
@@ -3461,7 +3464,7 @@ fn coinbase_outputs_are_decryptable_for_fake_v5_blocks() {
                 .find(|tx| tx.is_coinbase())
                 .expect("coinbase V5 tx");
 
-            let shielded_data = insert_fake_orchard_shielded_data(&mut transaction);
+            let shielded_data = insert_fake_v5_orchard_shielded_data(&mut transaction);
             shielded_data.flags = Flags::ENABLE_OUTPUTS;
 
             let action = fill_action_with_note_encryption_test_vector(
@@ -3494,7 +3497,7 @@ fn shielded_outputs_are_not_decryptable_for_fake_v5_blocks() {
                 .find(|tx| tx.is_coinbase())
                 .expect("V5 coinbase tx");
 
-            let shielded_data = insert_fake_orchard_shielded_data(&mut tx);
+            let shielded_data = insert_fake_v5_orchard_shielded_data(&mut tx);
             shielded_data.flags = Flags::ENABLE_OUTPUTS;
 
             let action = fill_action_with_note_encryption_test_vector(