From f56c1740cffb008305b435df6c5f868f76a08a23 Mon Sep 17 00:00:00 2001
From: keating <keating.dev@protonmail.com>
Date: Tue, 19 Mar 2024 13:25:21 -0400
Subject: [PATCH] Add invalidate nonce

---
 test/UniStaker.t.sol | 33 ++++++++++++++++++++++++++-------
 1 file changed, 26 insertions(+), 7 deletions(-)

diff --git a/test/UniStaker.t.sol b/test/UniStaker.t.sol
index 6f5a562..30a424e 100644
--- a/test/UniStaker.t.sol
+++ b/test/UniStaker.t.sol
@@ -2628,8 +2628,7 @@ contract SetAdmin is UniStakerTest {
 contract InvalidateNonce is UniStakerTest {
   using stdStorage for StdStorage;
 
-  // Calll increments nonce
-  function testFuzz_SucessfullyInvalidateNonceForMsgSender(address _caller, uint256 _initialNonce)
+  function testFuzz_SucessfullyInvalidatesNonceForMsgSender(address _caller, uint256 _initialNonce)
     public
   {
     vm.assume(_caller != address(0));
@@ -2648,8 +2647,7 @@ contract InvalidateNonce is UniStakerTest {
     assertEq(currentNonce, _initialNonce + 1, "Current nonce is incorrect");
   }
 
-  // Outside nonce owner cannot updateNonce
-  function testFuzz_NeedName(address _caller, address _nonceOwner, uint256 _initialNonce) public {
+  function testFuzz_NonOwnerCannotUpdateAnotherAddressesNonce(address _caller, address _nonceOwner, uint256 _initialNonce) public {
     vm.assume(_caller != address(0));
     vm.assume(_initialNonce != type(uint256).max);
 
@@ -2660,11 +2658,32 @@ contract InvalidateNonce is UniStakerTest {
 
     vm.prank(_caller);
     uniStaker.invalidateNonce();
+
+    uint256 currentNonce = uniStaker.nonces(_nonceOwner);
+
+	assertEq(currentNonce, _initialNonce);
   }
 
-  // It can be called multiple times
-  // It will invalidate StakeOnBehalf
-  // It will invalidate other stake on behalf methods
+  function testFuzz_NonceIncrementsWhenCalledMultipleTimes(address _caller, uint256, _initialNonce) public {
+    vm.assume(_caller != address(0));
+    vm.assume(_initialNonce != type(uint256).max);
+
+    // Set nonces in storage assert incremented by one
+    stdstore.target(address(uniStaker)).sig("nonces(address)").with_key(_caller).checked_write(
+      _initialNonce
+    );
+
+    vm.prank(_caller);
+    uniStaker.invalidateNonce();
+
+    vm.prank(_caller);
+    uniStaker.invalidateNonce();
+
+    uint256 currentNonce = uniStaker.nonces(_caller);
+
+    assertEq(currentNonce, _initialNonce + 2, "Current nonce is incorrect");
+
+  }
 }
 
 contract UniStakerRewardsTest is UniStakerTest {