From 66b33172e492f3b863157ac6adbd410c9a17751e Mon Sep 17 00:00:00 2001
From: Deirdre Connolly <deirdre@zfnd.org>
Date: Tue, 24 Mar 2020 20:01:01 -0400
Subject: [PATCH] Impl Arbitrary for SproutShieldedAddress and add roundtrip
 test

---
 zebra-chain/src/addresses/sprout.rs | 42 ++++++++++++++++++++++++++++-
 1 file changed, 41 insertions(+), 1 deletion(-)

diff --git a/zebra-chain/src/addresses/sprout.rs b/zebra-chain/src/addresses/sprout.rs
index 86ea92d89..47a7234d0 100644
--- a/zebra-chain/src/addresses/sprout.rs
+++ b/zebra-chain/src/addresses/sprout.rs
@@ -5,7 +5,7 @@ use std::{fmt, io};
 use bs58;
 
 #[cfg(test)]
-use proptest::{arbitrary::Arbitrary, collection::vec, prelude::*};
+use proptest::{arbitrary::Arbitrary, array, prelude::*};
 
 use crate::{
     keys::sprout,
@@ -93,3 +93,43 @@ impl ZcashDeserialize for SproutShieldedAddress {
         })
     }
 }
+
+#[cfg(test)]
+impl Arbitrary for SproutShieldedAddress {
+    type Parameters = ();
+
+    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
+        (
+            any::<Network>(),
+            array::uniform32(any::<u8>()),
+            array::uniform32(any::<u8>()),
+        )
+            .prop_map(|(network, paying_key_bytes, transmission_key_bytes)| {
+                return Self {
+                    network,
+                    paying_key: sprout::PayingKey(paying_key_bytes),
+                    transmission_key: sprout::TransmissionKey::from(transmission_key_bytes),
+                };
+            })
+            .boxed()
+    }
+
+    type Strategy = BoxedStrategy<Self>;
+}
+
+#[cfg(test)]
+proptest! {
+
+    #[test]
+    fn sprout_address_roundtrip(zaddr in any::<SproutShieldedAddress>()) {
+
+        let mut data = Vec::new();
+
+        zaddr.zcash_serialize(&mut data).expect("sprout z-addr should serialize");
+
+        let zaddr2 = SproutShieldedAddress::zcash_deserialize(&data[..])
+            .expect("randomized sprout z-addr should deserialize");
+
+        prop_assert_eq![zaddr, zaddr2];
+    }
+}