Add Arbitrary for Principal and extend random for Func

dfinity · Oct 27, 2023 · a968a82 · a968a82
1 parent 782c437
commit a968a82
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 1 deletion.
diff --git a/rust/candid/src/parser/random.rs b/rust/candid/src/parser/random.rs
@@ -140,7 +140,7 @@ impl<'a> GenState<'a> {
             TypeInner::Int64 => IDLValue::Int64(arbitrary_num(u, self.config.range)?),
             TypeInner::Float32 => IDLValue::Float32(u.arbitrary()?),
             TypeInner::Float64 => IDLValue::Float64(u.arbitrary()?),
-            TypeInner::Principal => IDLValue::Principal(crate::Principal::anonymous()),
+            TypeInner::Principal => IDLValue::Principal(crate::Principal::arbitrary(u)?),
             TypeInner::Text => {
                 IDLValue::Text(arbitrary_text(u, &self.config.text, &self.config.width)?)
             }
@@ -204,6 +204,9 @@ impl<'a> GenState<'a> {
                 };
                 IDLValue::Variant(VariantValue(Box::new(field), idx as u64))
             }
+            TypeInner::Func(_) => {
+                IDLValue::Func(crate::Principal::arbitrary(u)?, String::arbitrary(u)?)
+            }
             _ => unimplemented!(),
         });
         self.pop_state(old_config, ty, false);

diff --git a/rust/candid/src/types/principal.rs b/rust/candid/src/types/principal.rs
@@ -1,4 +1,6 @@
 use super::{CandidType, Serializer, Type, TypeInner};
+#[cfg(feature = "random")]
+use arbitrary::{Arbitrary, Result as ArbitraryResult, Unstructured};
 use serde::{Deserialize, Serialize};
 use sha2::{Digest, Sha224};
 use std::convert::TryFrom;
@@ -365,3 +367,28 @@ impl<'de> serde::Deserialize<'de> for Principal {
         }
     }
 }
+
+#[cfg(feature = "random")]
+impl<'a> Arbitrary<'a> for Principal {
+    fn arbitrary(u: &mut Unstructured<'a>) -> ArbitraryResult<Self> {
+        let principal = match u8::arbitrary(u)? {
+            u8::MAX => Principal::management_canister(),
+            254u8 => Principal::anonymous(),
+            _ => {
+                let length: usize = u.int_in_range(1..=Principal::MAX_LENGTH_IN_BYTES)?;
+                let mut result: Vec<u8> = Vec::with_capacity(length);
+                for _ in 0..length {
+                    result.push(u8::arbitrary(u)?);
+                }
+                // non-anonymous principal cannot have type ANONYMOUS
+                // adapt by changing the last byte.
+                let last = result.last_mut().unwrap();
+                if *last == 4_u8 {
+                    *last = u8::MAX
+                }
+                Principal::try_from(&result[..]).unwrap()
+            }
+        };
+        Ok(principal)
+    }
+}