From b9240971fd37e7a62ed1960182d87ab9901ecfd9 Mon Sep 17 00:00:00 2001
From: Deirdre Connolly <deirdre@zfnd.org>
Date: Thu, 12 Mar 2020 21:02:23 -0400
Subject: [PATCH] Add impl Arbitrary for TransparentAddress and a roundtrip
 proptest

addresses.rs is now large with just transparent stuff, i am already planning to break this up
further in the addresses/ module.
---
 zebra-chain/src/addresses.rs | 59 ++++++++++++++++++++++++++++++++++--
 zebra-chain/src/lib.rs       |  4 +++
 2 files changed, 61 insertions(+), 2 deletions(-)

diff --git a/zebra-chain/src/addresses.rs b/zebra-chain/src/addresses.rs
index 45a09e61b..b5b24c536 100644
--- a/zebra-chain/src/addresses.rs
+++ b/zebra-chain/src/addresses.rs
@@ -8,7 +8,7 @@ use secp256k1::PublicKey;
 use sha2::Sha256;
 
 #[cfg(test)]
-use proptest_derive::Arbitrary;
+use proptest::{arbitrary::Arbitrary, collection::vec, prelude::*};
 
 use crate::{
     serialization::{SerializationError, ZcashDeserialize, ZcashSerialize},
@@ -28,7 +28,6 @@ use crate::{
 ///
 /// https://zips.z.cash/protocol/protocol.pdf#transparentaddrencoding
 #[derive(Copy, Clone, Eq, PartialEq)]
-// #[cfg_attr(test, derive(Arbitrary))]
 pub enum TransparentAddress {
     /// P2SH (Pay to Script Hash) addresses
     PayToScriptHash {
@@ -158,6 +157,46 @@ impl ZcashDeserialize for TransparentAddress {
     }
 }
 
+#[cfg(test)]
+impl TransparentAddress {
+    fn p2pkh_strategy() -> impl Strategy<Value = Self> {
+        (any::<Network>(), vec(any::<u8>(), 20))
+            .prop_map(|(network, payload_bytes)| {
+                let mut bytes = [0; 20];
+                bytes.copy_from_slice(payload_bytes.as_slice());
+                return Self::PayToPublicKeyHash {
+                    network: network,
+                    pub_key_hash: bytes,
+                };
+            })
+            .boxed()
+    }
+
+    fn p2sh_strategy() -> impl Strategy<Value = Self> {
+        (any::<Network>(), vec(any::<u8>(), 20))
+            .prop_map(|(network, payload_bytes)| {
+                let mut bytes = [0; 20];
+                bytes.copy_from_slice(payload_bytes.as_slice());
+                return Self::PayToScriptHash {
+                    network: network,
+                    script_hash: bytes,
+                };
+            })
+            .boxed()
+    }
+}
+
+#[cfg(test)]
+impl Arbitrary for TransparentAddress {
+    type Parameters = ();
+
+    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
+        prop_oneof![Self::p2pkh_strategy(), Self::p2sh_strategy(),].boxed()
+    }
+
+    type Strategy = BoxedStrategy<Self>;
+}
+
 #[cfg(test)]
 mod tests {
 
@@ -195,3 +234,19 @@ mod tests {
         );
     }
 }
+
+#[cfg(test)]
+proptest! {
+
+    #[test]
+    fn transparent_address_roundtrip(taddr in any::<TransparentAddress>()) {
+
+        let mut data = Vec::new();
+
+        taddr.zcash_serialize(&mut data).expect("t-addr should serialize");
+
+        let taddr2 = TransparentAddress::zcash_deserialize(&data[..]).expect("randomized t-addr should deserialize");
+
+        prop_assert_eq![taddr, taddr2];
+    }
+}
diff --git a/zebra-chain/src/lib.rs b/zebra-chain/src/lib.rs
index 81141c2cf..bbba8c100 100644
--- a/zebra-chain/src/lib.rs
+++ b/zebra-chain/src/lib.rs
@@ -23,8 +23,12 @@ pub mod types;
 
 pub use redjubjub;
 
+#[cfg(test)]
+use proptest_derive::Arbitrary;
+
 /// An enum describing the possible network choices.
 #[derive(Copy, Clone, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)]
+#[cfg_attr(test, derive(Arbitrary))]
 pub enum Network {
     /// The production mainnet.
     Mainnet,