From ec53586b72757963e768e733392581c85d9a26e4 Mon Sep 17 00:00:00 2001
From: refcell <abigger87@gmail.com>
Date: Fri, 30 Aug 2024 09:05:22 -0400
Subject: [PATCH] feat(net): add arbitrary prop-based testing

---
 Cargo.lock                  | 31 +++++++++++++++++++++++++++++++
 Cargo.toml                  |  4 ++++
 crates/net/Cargo.toml       | 11 +++++++++++
 crates/net/src/types/enr.rs | 16 ++++++++++++++--
 4 files changed, 60 insertions(+), 2 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 66fed3b..e8f21a6 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -819,6 +819,24 @@ dependencies = [
  "syn 2.0.76",
 ]
 
+[[package]]
+name = "arbitrary"
+version = "1.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7d5a26814d8dcb93b0e5a0ff3c6d80a8843bafb21b39e8e18a6f05471870e110"
+dependencies = [
+ "derive_arbitrary",
+]
+
+[[package]]
+name = "arbtest"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "23909d5fb517fac2a8a4c887e847dbe41dd22ec46914586f5727980d0a193fdc"
+dependencies = [
+ "arbitrary",
+]
+
 [[package]]
 name = "ark-ff"
 version = "0.3.0"
@@ -2189,6 +2207,17 @@ dependencies = [
  "syn 1.0.109",
 ]
 
+[[package]]
+name = "derive_arbitrary"
+version = "1.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "67e77553c4162a157adbf834ebae5b415acbecbeafc7a74b0e886657506a7611"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.76",
+]
+
 [[package]]
 name = "derive_more"
 version = "0.99.18"
@@ -5192,6 +5221,8 @@ version = "0.0.0"
 dependencies = [
  "alloy",
  "alloy-rlp",
+ "arbitrary",
+ "arbtest",
  "discv5",
  "eyre",
  "futures",
diff --git a/Cargo.toml b/Cargo.toml
index e88ce6c..5e96db0 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -166,6 +166,10 @@ openssl = { version = "0.10.66", features = ["vendored"] }
 libp2p-identity = { version = "0.2.9", features = [ "secp256k1" ] }
 libp2p = { version = "0.54.0", features = ["macros", "tokio", "tcp", "noise", "gossipsub", "ping", "yamux"] }
 
+# Testing
+arbtest = "0.3"
+arbitrary = { version = "1", features = ["derive"] }
+
 # Misc
 tracing = "0.1.0"
 tracing-subscriber = "0.3.18"
diff --git a/crates/net/Cargo.toml b/crates/net/Cargo.toml
index 7958beb..f25b554 100644
--- a/crates/net/Cargo.toml
+++ b/crates/net/Cargo.toml
@@ -33,3 +33,14 @@ tokio.workspace = true
 tracing.workspace = true
 lazy_static.workspace = true
 unsigned-varint.workspace = true
+
+# `arbitrary` feature dependencies
+arbitrary = { workspace = true, optional = true }
+
+[dev-dependencies]
+arbtest.workspace = true
+arbitrary.workspace = true
+
+[features]
+default = []
+arbitrary = ["dep:arbitrary"]
diff --git a/crates/net/src/types/enr.rs b/crates/net/src/types/enr.rs
index 7c972c6..71a0612 100644
--- a/crates/net/src/types/enr.rs
+++ b/crates/net/src/types/enr.rs
@@ -9,6 +9,7 @@ pub const OP_CL_KEY: &str = "opstack";
 
 /// The unique L2 network identifier
 #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
+#[cfg_attr(any(test, feature = "arbitrary"), derive(arbitrary::Arbitrary))]
 pub struct OpStackEnr {
     /// Chain ID
     pub chain_id: u64,
@@ -64,7 +65,18 @@ mod tests {
     use alloy::primitives::{bytes, Bytes};
 
     #[test]
-    fn test_encode_decode_op_enr() {
+    fn rountrip_op_stack_enr() {
+        arbtest::arbtest(|u| {
+            let op_stack_enr = OpStackEnr::new(u.arbitrary()?, 0);
+            let bytes = alloy_rlp::encode(op_stack_enr).to_vec();
+            let decoded = OpStackEnr::decode(&mut &bytes[..]).unwrap();
+            assert_eq!(decoded, op_stack_enr);
+            Ok(())
+        });
+    }
+
+    #[test]
+    fn test_op_mainnet_enr() {
         let op_enr = OpStackEnr::new(10, 0);
         let bytes = alloy_rlp::encode(op_enr).to_vec();
         assert_eq!(Bytes::from(bytes.clone()), bytes!("820A00"));
@@ -73,7 +85,7 @@ mod tests {
     }
 
     #[test]
-    fn test_encode_decode_base_enr() {
+    fn test_base_mainnet_enr() {
         let base_enr = OpStackEnr::new(8453, 0);
         let bytes = alloy_rlp::encode(base_enr).to_vec();
         assert_eq!(Bytes::from(bytes.clone()), bytes!("83854200"));