refactor(test): more rigorously use proptest framework

TritonVM · Jan 23, 2024 · a27ca6d · a27ca6d
1 parent d16fc4a
commit a27ca6d
Showing 1 changed file with 13 additions and 27 deletions.
diff --git a/triton-vm/src/proof.rs b/triton-vm/src/proof.rs
@@ -74,59 +74,45 @@ impl Claim {
 mod tests {
     use assert2::assert;
     use proptest::collection::vec;
+    use proptest::prelude::*;
     use proptest_arbitrary_interop::arb;
-    use rand::random;
     use test_strategy::proptest;
     use twenty_first::shared_math::b_field_element::BFieldElement;
     use twenty_first::shared_math::bfield_codec::BFieldCodec;
-    use twenty_first::shared_math::other::random_elements;
 
     use crate::proof_item::ProofItem;
     use crate::stark::StarkHasher;
 
     use super::*;
 
-    #[test]
-    fn decode_proof() {
-        let data: Vec<BFieldElement> = random_elements(348);
-        let proof = Proof(data);
-
+    #[proptest]
+    fn decode_proof(#[strategy(arb())] proof: Proof) {
         let encoded = proof.encode();
         let decoded = *Proof::decode(&encoded).unwrap();
-
-        assert!(proof == decoded);
+        prop_assert_eq!(proof, decoded);
     }
 
-    #[test]
-    fn decode_claim() {
-        let claim = Claim {
-            program_digest: random(),
-            input: random_elements(346),
-            output: random_elements(125),
-        };
-
+    #[proptest]
+    fn decode_claim(#[strategy(arb())] claim: Claim) {
         let encoded = claim.encode();
         let decoded = *Claim::decode(&encoded).unwrap();
-
-        assert!(claim.program_digest == decoded.program_digest);
-        assert!(claim.input == decoded.input);
-        assert!(claim.output == decoded.output);
+        prop_assert_eq!(claim, decoded);
     }
 
-    #[test]
-    fn proof_with_no_log_2_padded_height_gives_err() {
+    #[proptest(cases = 10)]
+    fn proof_with_no_log_2_padded_height_gives_err(#[strategy(arb())] root: Digest) {
         let mut proof_stream = ProofStream::<StarkHasher>::new();
-        proof_stream.enqueue(ProofItem::MerkleRoot(random()));
+        proof_stream.enqueue(ProofItem::MerkleRoot(root));
         let proof: Proof = proof_stream.into();
         let maybe_padded_height = proof.padded_height();
         assert!(maybe_padded_height.is_err());
     }
 
-    #[test]
-    fn proof_with_multiple_log_2_padded_height_gives_err() {
+    #[proptest(cases = 10)]
+    fn proof_with_multiple_log_2_padded_height_gives_err(#[strategy(arb())] root: Digest) {
         let mut proof_stream = ProofStream::<StarkHasher>::new();
         proof_stream.enqueue(ProofItem::Log2PaddedHeight(8));
-        proof_stream.enqueue(ProofItem::MerkleRoot(random()));
+        proof_stream.enqueue(ProofItem::MerkleRoot(root));
         proof_stream.enqueue(ProofItem::Log2PaddedHeight(7));
         let proof: Proof = proof_stream.into();
         let maybe_padded_height = proof.padded_height();