fix(TwabRewards): fix getCurrentEpochId edge cases

pooltogether · Dec 15, 2021 · e6081b9 · e6081b9
1 parent c6febd7
commit e6081b9
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 7 deletions.
diff --git a/contracts/TwabRewards.sol b/contracts/TwabRewards.sol
@@ -268,8 +268,7 @@ contract TwabRewards is ITwabRewards {
         unchecked {
             // promotionEndTimestamp >= block.timestamp
             require(
-                (_promotion.startTimestamp +
-                    (_promotion.epochDuration * _promotion.numberOfEpochs)) >= block.timestamp,
+                _getPromotionEndTimestamp(_promotion) >= block.timestamp,
                 "TwabRewards/promotion-inactive"
             );
         }
@@ -287,14 +286,36 @@ contract TwabRewards is ITwabRewards {
         return _promotion;
     }
 
+    /**
+        @notice Compute promotion end timestamp.
+        @param _promotion Promotion to compute end timestamp for
+        @return Promotion end timestamp
+    */
+    function _getPromotionEndTimestamp(Promotion memory _promotion)
+        internal
+        pure
+        returns (uint256)
+    {
+        return _promotion.startTimestamp + (_promotion.epochDuration * _promotion.numberOfEpochs);
+    }
+
     /**
         @notice Get the current epoch id of a promotion.
         @dev Epoch ids and their boolean values are tightly packed and stored in a uint256, so epoch id starts at 0.
-        @dev We calculate the current epoch id if the promotion has started. Otherwise, we return 0.
+        @dev We return the current epoch id if the promotion has not ended.
+        Otherwise, we return the last epoch id or 0 if the promotion has not started.
         @param _promotion Promotion to get current epoch for
         @return Epoch id
      */
     function _getCurrentEpochId(Promotion memory _promotion) internal view returns (uint256) {
+        if (block.timestamp > _getPromotionEndTimestamp(_promotion)) {
+            if (_promotion.numberOfEpochs > 0) {
+                unchecked {
+                    return _promotion.numberOfEpochs - uint8(1);
+                }
+            }
+        }
+
         if (block.timestamp > _promotion.startTimestamp) {
             unchecked {
                 // elapsedTimestamp / epochDurationTimestamp

diff --git a/test/TwabRewards.test.ts b/test/TwabRewards.test.ts
@@ -229,9 +229,16 @@ describe('TwabRewards', () => {
 
                 expect(await rewardToken.balanceOf(wallet1.address)).to.equal(transferredAmount);
 
-                expect(
-                    (await twabRewards.callStatic.getPromotion(promotionId)).numberOfEpochs,
-                ).to.equal(await twabRewards.callStatic.getCurrentEpochId(promotionId));
+                let latestEpochId = (await twabRewards.callStatic.getPromotion(promotionId))
+                    .numberOfEpochs;
+
+                if (latestEpochId !== 0) {
+                    latestEpochId--;
+                }
+
+                expect(latestEpochId).to.equal(
+                    await twabRewards.callStatic.getCurrentEpochId(promotionId),
+                );
 
                 // We burn tokens from wallet1 to reset balance
                 await rewardToken.burn(wallet1.address, transferredAmount);
@@ -297,8 +304,9 @@ describe('TwabRewards', () => {
                 .withArgs(promotionId, wallet1.address, transferredAmount);
 
             expect(await rewardToken.balanceOf(wallet1.address)).to.equal(transferredAmount);
+
             expect(
-                (await twabRewards.callStatic.getPromotion(promotionId)).numberOfEpochs,
+                (await twabRewards.callStatic.getPromotion(promotionId)).numberOfEpochs - 1,
             ).to.equal(await twabRewards.callStatic.getCurrentEpochId(promotionId));
 
             await expect(twabRewards.claimRewards(wallet2.address, promotionId, epochIds))
@@ -457,6 +465,30 @@ describe('TwabRewards', () => {
             expect(await twabRewards.callStatic.getCurrentEpochId(1)).to.equal(3);
         });
 
+        it('should return the first epoch id if the promotion has not started yet', async () => {
+            const startTimestamp = (await ethers.provider.getBlock('latest')).timestamp + 60;
+
+            await createPromotion(
+                ticket.address,
+                rewardToken,
+                tokensPerEpoch,
+                epochDuration,
+                numberOfEpochs,
+                startTimestamp,
+            ),
+
+            expect(await twabRewards.callStatic.getCurrentEpochId(1)).to.equal(0);
+        });
+
+        it('should return the last active epoch id if the promotion is inactive', async () => {
+            await createPromotion();
+            await increaseTime(epochDuration * 12);
+
+            expect(await twabRewards.callStatic.getCurrentEpochId(1)).to.equal(
+                numberOfEpochs - 1,
+            );
+        });
+
         it('should revert if promotion id passed is inexistent', async () => {
             await expect(twabRewards.callStatic.getCurrentEpochId(1)).to.be.revertedWith(
                 'TwabRewards/invalid-promotion',