From 8693c46c57b4f87e15158c55163eacf16c6f0682 Mon Sep 17 00:00:00 2001
From: xavikh <xavikh@gmail.com>
Date: Mon, 16 Dec 2024 23:22:25 +0000
Subject: [PATCH] Add some tests to resetVotes

---
 test/escrow/escrow/EscrowWithdraw.t.sol | 52 ++++++++++++++++++++++
 test/voting/GaugeVote.t.sol             | 58 +++++++++++++++++++++++--
 2 files changed, 107 insertions(+), 3 deletions(-)

diff --git a/test/escrow/escrow/EscrowWithdraw.t.sol b/test/escrow/escrow/EscrowWithdraw.t.sol
index 646e7c1..77d3856 100644
--- a/test/escrow/escrow/EscrowWithdraw.t.sol
+++ b/test/escrow/escrow/EscrowWithdraw.t.sol
@@ -224,6 +224,58 @@ contract TestWithdraw is EscrowBase, IEscrowCurveTokenStorage, IGaugeVote, ITick
         assertEq(nftLock.balanceOf(address(escrow)), 0);
         assertEq(escrow.totalLocked(), 0);
     }
+
+    function test_CanBeginWithdrawalDuringDistributionPeriod() public {
+        address _who = address(1);
+        uint128 _dep = 100e18;
+
+        // voting window is ea. 2 weeks + 1 hour
+        vm.warp(2 weeks + 1 hours + 1);
+
+        token.mint(_who, _dep);
+        uint tokenId;
+        vm.startPrank(_who);
+        {
+            token.approve(address(escrow), _dep);
+            tokenId = escrow.createLock(_dep);
+
+            // voting active after cooldown
+            // +1 week: voting ends
+            // +2 weeks: next voting period opens
+            vm.warp(block.timestamp + 2 weeks);
+
+            // make a vote
+            voter.vote(tokenId, votes);
+
+            // warp so cooldown crosses the week boundary
+            // and distribution period starts
+            vm.warp(block.timestamp + clock.epochNextCheckpointIn() + 1);
+            assertFalse(voter.votingActive());
+
+            nftLock.approve(address(escrow), tokenId);
+            escrow.resetVotesAndBeginWithdrawal(tokenId);
+        }
+        vm.stopPrank();
+
+        // must wait till after end of cooldown
+        vm.warp(block.timestamp + clock.epochNextCheckpointIn() + 1);
+
+        uint fee = queue.calculateFee(tokenId);
+
+        // withdraw
+        vm.prank(_who);
+        vm.expectEmit(true, true, false, true);
+        emit Withdraw(_who, tokenId, _dep - fee, block.timestamp, 0);
+        escrow.withdraw(tokenId);
+
+        // asserts
+        assertEq(token.balanceOf(address(queue)), fee);
+        assertEq(token.balanceOf(_who), _dep - fee);
+        assertEq(nftLock.balanceOf(_who), 0);
+        assertEq(nftLock.balanceOf(address(escrow)), 0);
+        assertEq(escrow.totalLocked(), 0);
+    }
+
     // HAL-13: locks are re-used causing reverts and duplications
     function testCanCreateLockAfterBurning() public {
         address USER1 = address(1);
diff --git a/test/voting/GaugeVote.t.sol b/test/voting/GaugeVote.t.sol
index aa51d0e..ed3caf2 100644
--- a/test/voting/GaugeVote.t.sol
+++ b/test/voting/GaugeVote.t.sol
@@ -66,9 +66,9 @@ contract TestGaugeVote is GaugeVotingBase {
         voter.createGauge(gauge, "metadata");
     }
 
-    function testFuzz_cannotVoteOutsideVotingWindow(uint256 time) public {
+    function testFuzz_cannotVoteOutsideVotingWindow(uint256 _time) public {
         // warp to a random time
-        vm.warp(time);
+        vm.warp(_time);
 
         // should now be inactive (we don't test this part herewe have the epoch logic tests)
         vm.assume(!voter.votingActive());
@@ -97,7 +97,7 @@ contract TestGaugeVote is GaugeVotingBase {
         assertEq(voter.isVoting(tokenId), true);
 
         // warp to the next distribution period
-        vm.warp(block.timestamp + 1 weeks);
+        _increaseTime(1 weeks);
         vm.assume(!voter.votingActive());
 
         // try to reset
@@ -107,7 +107,59 @@ contract TestGaugeVote is GaugeVotingBase {
         }
         vm.stopPrank();
 
+        // check the vote
         assertEq(voter.isVoting(tokenId), false);
+        assertEq(voter.gaugesVotedFor(tokenId).length, 0);
+        assertEq(voter.votes(tokenId, gauge), 0);
+        assertEq(voter.usedVotingPower(tokenId), 0);
+
+        // global state
+        assertEq(voter.totalVotingPowerCast(), 0);
+        assertEq(voter.gaugeVotes(gauge), 0);
+    }
+
+    function testFuzz_canResetAnytime(uint _time) public {
+        // create the vote
+        votes.push(GaugeVote(1, gauge));
+
+        // vote
+        vm.startPrank(owner);
+        {
+            voter.vote(tokenId, votes);
+        }
+        vm.stopPrank();
+
+        // check the vote
+        uint newVotingPower = escrow.votingPower(tokenId);
+        assertEq(voter.isVoting(tokenId), true);
+        assertEq(voter.gaugesVotedFor(tokenId).length, 1);
+        assertEq(voter.gaugesVotedFor(tokenId)[0], gauge);
+        assertEq(voter.votes(tokenId, gauge), newVotingPower);
+        assertEq(voter.usedVotingPower(tokenId), newVotingPower);
+
+        // global state
+        assertEq(voter.totalVotingPowerCast(), newVotingPower);
+        assertEq(voter.gaugeVotes(gauge), newVotingPower);
+
+        // warp to the next distribution period
+        _increaseTime(_time);
+
+        // try to reset
+        vm.startPrank(owner);
+        {
+            voter.reset(tokenId);
+        }
+        vm.stopPrank();
+
+        // check the vote
+        assertEq(voter.isVoting(tokenId), false);
+        assertEq(voter.gaugesVotedFor(tokenId).length, 0);
+        assertEq(voter.votes(tokenId, gauge), 0);
+        assertEq(voter.usedVotingPower(tokenId), 0);
+
+        // global state
+        assertEq(voter.totalVotingPowerCast(), 0);
+        assertEq(voter.gaugeVotes(gauge), 0);
     }
 
     // can't vote if you don't own the token