Update VecArb to support Principals

demergent-labs · Oct 24, 2023 · 6d6c058 · 6d6c058
1 parent d4fdbec
commit 6d6c058
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 6 deletions.
diff --git a/property_tests/arbitraries/candid/constructed/vec_arb.ts b/property_tests/arbitraries/candid/constructed/vec_arb.ts
@@ -1,3 +1,4 @@
+import { Principal } from '@dfinity/principal';
 import fc from 'fast-check';
 import { IntArb } from '../primitive/ints/int_arb';
 import { Int8Arb } from '../primitive/ints/int8_arb';
@@ -9,6 +10,7 @@ import { Nat8Arb } from '../primitive/nats/nat8_arb';
 import { Nat16Arb } from '../primitive/nats/nat16_arb';
 import { Nat32Arb } from '../primitive/nats/nat32_arb';
 import { Nat64Arb } from '../primitive/nats/nat64_arb';
+import { PrincipalArb } from '../reference/principal_arb';
 
 // TODO look into making this recursive
 // TODO we want to be able to have vecs of vecs
@@ -35,12 +37,26 @@ export const VecArb = fc.oneof(
         .map((sample) => createVecArbWrapper(sample, 'Vec(nat32)')),
     fc
         .array(Nat64Arb)
-        .map((sample) => createVecArbWrapper(sample, 'Vec(nat64)'))
+        .map((sample) => createVecArbWrapper(sample, 'Vec(nat64)')),
+    fc
+        .array(PrincipalArb)
+        .map((sample) =>
+            createVecArbWrapper(
+                sample,
+                'Vec(Principal)',
+                (a: Principal, b: Principal) => a.toText() === b.toText()
+            )
+        )
 );
 
-function createVecArbWrapper(sample: any[], candidType: string) {
+function createVecArbWrapper(
+    sample: any[],
+    candidType: string,
+    equalityCheck: (a: any, b: any) => boolean = (a, b) => a === b
+) {
     return {
         vec: sample,
-        candidType
+        candidType,
+        equalityCheck
     };
 }
diff --git a/property_tests/tests/vec/test/test.ts b/property_tests/tests/vec/test/test.ts
@@ -34,6 +34,9 @@ const VecTestArb = fc
 
         const expectedResult = vecWrappers[0]?.vec ?? [];
 
+        const equalityCheck =
+            vecWrappers[0]?.equalityCheck ?? ((a, b) => a === b);
+
         return {
             functionName,
             imports: [
@@ -47,6 +50,7 @@ const VecTestArb = fc
                 'nat16',
                 'nat32',
                 'nat64',
+                'Principal',
                 'Vec'
             ],
             paramCandidTypes: paramCandidTypes.join(', '),
@@ -69,7 +73,11 @@ const VecTestArb = fc
                     );
 
                     return {
-                        Ok: primitiveArraysAreEqual(result, expectedResult)
+                        Ok: primitiveArraysAreEqual(
+                            result,
+                            expectedResult,
+                            equalityCheck
+                        )
                     };
                 }
             }
@@ -78,15 +86,19 @@ const VecTestArb = fc
 
 runPropTests(VecTestArb);
 
-function primitiveArraysAreEqual(arr1: any, arr2: any) {
+function primitiveArraysAreEqual(
+    arr1: any,
+    arr2: any,
+    equalityCheck: (a: any, b: any) => boolean
+) {
     // Check if both arrays have the same length
     if (arr1.length !== arr2.length) {
         return false;
     }
 
     // Loop through each element to check for equality
     for (let i = 0; i < arr1.length; i++) {
-        if (arr1[i] !== arr2[i]) {
+        if (!equalityCheck(arr1[i], arr2[i])) {
             return false;
         }
     }