From a37c0460df64df8f3749582320d6a3842e662071 Mon Sep 17 00:00:00 2001
From: Tal Derei <talderei99@gmail.com>
Date: Fri, 13 Dec 2024 11:25:31 -0800
Subject: [PATCH] testing: add proptest and filter zero amounts in fallible
 conversion

---
 crates/core/asset/src/balance.rs | 41 +++++++++++++++++++++++++++++---
 1 file changed, 38 insertions(+), 3 deletions(-)

diff --git a/crates/core/asset/src/balance.rs b/crates/core/asset/src/balance.rs
index ecdea31da7..f5013f1d83 100644
--- a/crates/core/asset/src/balance.rs
+++ b/crates/core/asset/src/balance.rs
@@ -61,8 +61,10 @@ impl TryFrom<pb::Balance> for Balance {
                 .value
                 .ok_or_else(|| anyhow::anyhow!("missing value"))?;
             let value: Value = proto_value.try_into()?;
-            let amount = NonZeroU128::new(value.amount.into())
-                .ok_or_else(|| anyhow::anyhow!("amount is zero"))?;
+            let amount = match NonZeroU128::new(value.amount.into()) {
+                Some(amount) => amount,
+                None => continue,
+            };
 
             // The 'negated' flag in SignedValue determines the imbalance type:
             // true = Required, false = Provided
@@ -539,7 +541,10 @@ impl std::ops::Add for BalanceVar {
 
 #[cfg(test)]
 mod test {
-    use crate::{asset::Metadata, STAKING_TOKEN_ASSET_ID};
+    use crate::{
+        asset::{self, Metadata},
+        STAKING_TOKEN_ASSET_ID,
+    };
     use ark_ff::Zero;
     use decaf377::Fr;
     use once_cell::sync::Lazy;
@@ -1023,4 +1028,34 @@ mod test {
             Some(proto_amount)
         );
     }
+
+    fn test_balance_serialization_round_tripping_example(
+        parts: Vec<(u8, i64)>,
+    ) -> anyhow::Result<()> {
+        let proto = pb::Balance {
+            values: parts
+                .into_iter()
+                .map(|(asset, amount)| pb::balance::SignedValue {
+                    value: Some(pb::Value {
+                        amount: Some(Amount::from(amount.abs() as u64).into()),
+                        asset_id: Some(asset::Id::try_from([asset; 32]).unwrap().into()),
+                    }),
+                    negated: amount < 0,
+                })
+                .collect(),
+        };
+        let p_to_d = Balance::try_from(proto)?;
+        let p_to_d_to_p = pb::Balance::from(p_to_d.clone());
+        let p_to_d_to_p_to_d = Balance::try_from(p_to_d_to_p)?;
+        assert_eq!(p_to_d, p_to_d_to_p_to_d);
+        Ok(())
+    }
+
+    proptest! {
+        #[test]
+        fn test_balance_serialization_roundtripping(parts in proptest::collection::vec(((0u8..16u8), any::<i64>()), 0..100)) {
+            // To get better errors
+            assert!(matches!(test_balance_serialization_round_tripping_example(parts), Ok(_)));
+        }
+    }
 }