add feature flags to arbitrary_with

paradigmxyz · Oct 4, 2023 · 109ca9e · 109ca9e
1 parent 76742a4
commit 109ca9e
Showing 1 changed file with 11 additions and 3 deletions.
diff --git a/crates/primitives/src/transaction/tx_value.rs b/crates/primitives/src/transaction/tx_value.rs
@@ -135,9 +135,17 @@ impl<'a> arbitrary::Arbitrary<'a> for TxValue {
 impl proptest::arbitrary::Arbitrary for TxValue {
     type Parameters = ParamsFor<()>;
     fn arbitrary_with(_: Self::Parameters) -> Self::Strategy {
-        proptest::prelude::any::<u128>()
-            .prop_map(|num| Self::try_from(num).expect("to fit"))
-            .boxed()
+        #[cfg(feature = "value-256")]
+        {
+            proptest::prelude::any::<U256>().prop_map(Self).boxed()
+        }
+
+        #[cfg(not(feature = "value-256"))]
+        {
+            proptest::prelude::any::<u128>()
+                .prop_map(|num| Self::try_from(num).expect("to fit"))
+                .boxed()
+        }
     }
 
     type Strategy = BoxedStrategy<TxValue>;