From 9e8ba3565143074c7168dd9f75e491234b52c72b Mon Sep 17 00:00:00 2001
From: Roman <zajic@zajic.net>
Date: Mon, 29 Sep 2025 16:16:16 +1000
Subject: [PATCH] test: erasure with fullErase idCommitments

---
 test/WakuRlnV2.fuzz.t.sol | 65 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 65 insertions(+)

diff --git a/test/WakuRlnV2.fuzz.t.sol b/test/WakuRlnV2.fuzz.t.sol
index 52464b6..b10f217 100644
--- a/test/WakuRlnV2.fuzz.t.sol
+++ b/test/WakuRlnV2.fuzz.t.sol
@@ -68,4 +68,69 @@ contract WakuRlnV2Test is Test {
         }
         assertEq(w.currentTotalRateLimit(), totalExpected);
     }
+
+    // Helper function to register a single membership (reusable in tests)
+    function _registerMembership(uint256 idCommitment, uint32 rateLimit) internal {
+        (, uint256 price) = w.priceCalculator().calculate(rateLimit);
+        token.approve(address(w), price);
+        w.register(idCommitment, rateLimit, new uint256[](0));
+    }
+
+    // Fuzz Test: Erasure with Random IDs and Time Deltas
+    function testFuzz_Erasure(bool fullErase, uint8 subsetMask) external {
+        vm.assume(subsetMask > 0 && subsetMask < 16);
+        // Setup: Register multiple memberships to allow fuzzing various IDs
+        uint32 rateLimit = w.minMembershipRateLimit();
+        uint256 initialTotal = 0;
+        for (uint256 i = 1; i <= 4; i++) {
+            // Register up to 4 for <5 constraint
+            if (w.currentTotalRateLimit() + rateLimit <= w.maxTotalRateLimit()) {
+                _registerMembership(i, rateLimit);
+                initialTotal += rateLimit;
+            }
+        }
+
+        // Fuzz time warp (constrain to > active + grace for expiration, and < 1180 year to reduce range and speed up)
+        uint256 minDelta =
+            uint256(w.activeDurationForNewMemberships()) + uint256(w.gracePeriodDurationForNewMemberships()) + 1;
+        vm.warp(block.timestamp + minDelta);
+
+        // Constrain array
+        uint256 len = 0;
+        if (subsetMask & 1 != 0) len++;
+        if (subsetMask & 2 != 0) len++;
+        if (subsetMask & 4 != 0) len++;
+        if (subsetMask & 8 != 0) len++;
+        uint256[] memory idCommitments = new uint256[](len);
+        uint256 idx = 0;
+        if (subsetMask & 1 != 0) idCommitments[idx++] = 1;
+        if (subsetMask & 2 != 0) idCommitments[idx++] = 2;
+        if (subsetMask & 4 != 0) idCommitments[idx++] = 3;
+        if (subsetMask & 8 != 0) idCommitments[idx++] = 4;
+
+        // Record indices before erasure
+        uint32[] memory indices = new uint32[](idCommitments.length);
+        for (uint256 j = 0; j < idCommitments.length; j++) {
+            (, indices[j],) = w.getMembershipInfo(idCommitments[j]); // Get original index
+        }
+
+        w.eraseMemberships(idCommitments, fullErase);
+
+        // Assert invariants: For each ID, check erased if conditions met
+        uint256 erasedTotal = 0;
+        uint256 rateLimitCast = uint256(rateLimit);
+        for (uint256 j = 0; j < idCommitments.length; j++) {
+            assertFalse(w.isInMembershipSet(idCommitments[j]));
+            (uint32 rl,, uint256 commitment) = w.getMembershipInfo(idCommitments[j]);
+            assertEq(rl, 0);
+            assertEq(commitment, 0); // Info returns 0 if erased
+            if (indices[j] < w.nextFreeIndex()) {
+                // Valid index
+                uint256 expectedCommitment = fullErase ? 0 : PoseidonT3.hash([idCommitments[j], rateLimitCast]);
+                assertEq(w.getRateCommitmentsInRangeBoundsInclusive(indices[j], indices[j])[0], expectedCommitment);
+            }
+            erasedTotal += rateLimit; // Assuming all were valid to erase
+        }
+        assertEq(w.currentTotalRateLimit(), initialTotal - erasedTotal);
+    }
 }