stCelo balance (#205)

celo-org · Jul 24, 2024 · cfe8b13 · cfe8b13
1 parent a05d7b1
commit cfe8b13
Show file tree

Hide file tree

Showing 5 changed files with 99 additions and 11 deletions.
diff --git a/contracts/DefaultStrategy.sol b/contracts/DefaultStrategy.sol
@@ -394,8 +394,8 @@ contract DefaultStrategy is Errors, UUPSOwnableUpgradeable, Managed, Pausable {
             expectedToStCelo - actualToStCelo
         );
 
-        updateGroupStCelo(fromGroup, toMove, false);
-        updateGroupStCelo(toGroup, toMove, true);
+        _updateGroupStCelo(fromGroup, toMove, false);
+        _updateGroupStCelo(toGroup, toMove, true);
 
         trySort(fromGroup, stCeloInGroup[fromGroup], false);
         trySort(toGroup, stCeloInGroup[toGroup], true);
@@ -436,7 +436,7 @@ contract DefaultStrategy is Errors, UUPSOwnableUpgradeable, Managed, Pausable {
 
             groups[groupsIndex] = votedGroup;
             celoAmount -= votes[groupsIndex];
-            updateGroupStCelo(
+            _updateGroupStCelo(
                 votedGroup,
                 IManager(manager).toStakedCelo(votes[groupsIndex]),
                 false
@@ -508,6 +508,20 @@ contract DefaultStrategy is Errors, UUPSOwnableUpgradeable, Managed, Pausable {
         activatableGroups.add(group);
     }
 
+    /**
+     * @notice Updates the total stCELO in default group strategy.
+     * @param group The group address.
+     * @param stCeloAmount The amount of stCELO.
+     * @param add Whether to add or subtract.
+     */
+    function updateGroupStCelo(
+        address group,
+        uint256 stCeloAmount,
+        bool add
+    ) external onlyOwner {
+        _updateGroupStCelo(group, stCeloAmount, add);
+    }
+
     /**
      * @notice Returns the number of active groups.
      * @return The number of active groups.
@@ -644,7 +658,7 @@ contract DefaultStrategy is Errors, UUPSOwnableUpgradeable, Managed, Pausable {
      * @param stCeloAmount The amount of stCELO.
      * @param add Whether to add or substract.
      */
-    function updateGroupStCelo(
+    function _updateGroupStCelo(
         address group,
         uint256 stCeloAmount,
         bool add
@@ -677,7 +691,7 @@ contract DefaultStrategy is Errors, UUPSOwnableUpgradeable, Managed, Pausable {
         uint256 groupTotalStCeloVotes = stCeloInGroup[group];
 
         if (groupTotalStCeloVotes > 0) {
-            updateGroupStCelo(group, groupTotalStCeloVotes, false);
+            _updateGroupStCelo(group, groupTotalStCeloVotes, false);
             _generateDepositVoteDistribution(
                 IManager(manager).toCelo(groupTotalStCeloVotes),
                 address(0)
@@ -721,7 +735,11 @@ contract DefaultStrategy is Errors, UUPSOwnableUpgradeable, Managed, Pausable {
             votes[groupsIndex] = Math.min(receivableVotes, celoAmount);
             groups[groupsIndex] = votedGroup;
             celoAmount -= votes[groupsIndex];
-            updateGroupStCelo(votedGroup, IManager(manager).toStakedCelo(votes[groupsIndex]), true);
+            _updateGroupStCelo(
+                votedGroup,
+                IManager(manager).toStakedCelo(votes[groupsIndex]),
+                true
+            );
             trySort(votedGroup, stCeloInGroup[votedGroup], true);
 
             if (sorted) {

diff --git a/contracts/SpecificGroupStrategy.sol b/contracts/SpecificGroupStrategy.sol
@@ -229,6 +229,20 @@ contract SpecificGroupStrategy is Errors, UUPSOwnableUpgradeable, Managed, Pausa
         _blockGroup(group);
     }
 
+    /**
+     * @notice Updates the total stCELO in specific group strategy.
+     * @param group The group address.
+     * @param stCeloAmount The amount of stCELO.
+     * @param add Whether to add or subtract.
+     */
+    function updateGroupStCelo(
+        address group,
+        uint256 stCeloAmount,
+        bool add
+    ) external onlyOwner {
+        _updateGroupStCelo(group, stCeloAmount, add);
+    }
+
     /**
      * @notice Used to withdraw CELO from a specific group
      * that account voted for previously. It is expected that strategy will be balanced.
@@ -253,7 +267,7 @@ contract SpecificGroupStrategy is Errors, UUPSOwnableUpgradeable, Managed, Pausa
             revert CantWithdrawAccordingToStrategy(group);
         }
 
-        updateGroupStCelo(group, stCeloWithdrawalAmount, false);
+        _updateGroupStCelo(group, stCeloWithdrawalAmount, false);
 
         uint256 overflowingStCelo = stCeloInGroupOverflowed[group];
         uint256 unhealthyStCelo = stCeloInGroupUnhealthy[group];
@@ -311,7 +325,7 @@ contract SpecificGroupStrategy is Errors, UUPSOwnableUpgradeable, Managed, Pausa
         uint256 stCeloAmount
     ) external onlyManager returns (address[] memory finalGroups, uint256[] memory finalVotes) {
         votedGroups.add(group);
-        updateGroupStCelo(group, stCeloAmount, true);
+        _updateGroupStCelo(group, stCeloAmount, true);
 
         if (groupHealth.isGroupValid(group) && !blockedGroups.contains(group)) {
             uint256 receivableVotes = IManager(manager).getReceivableVotesForGroup(group);
@@ -509,7 +523,7 @@ contract SpecificGroupStrategy is Errors, UUPSOwnableUpgradeable, Managed, Pausa
      * @param stCeloAmount The amount of stCELO.
      * @param add Whether to add or substract.
      */
-    function updateGroupStCelo(
+    function _updateGroupStCelo(
         address group,
         uint256 stCeloAmount,
         bool add

diff --git a/contracts/mock/MockDefaultStrategy.sol b/contracts/mock/MockDefaultStrategy.sol
@@ -8,12 +8,12 @@ contract MockDefaultStrategy is DefaultStrategy {
     receive() external payable {}
 
     function addToStrategyTotalStCeloVotesPublic(address strategy, uint256 stCeloAmount) public {
-        updateGroupStCelo(strategy, stCeloAmount, true);
+        _updateGroupStCelo(strategy, stCeloAmount, true);
     }
 
     function subtractFromStrategyTotalStCeloVotesPublic(address strategy, uint256 stCeloAmount)
         internal
     {
-        updateGroupStCelo(strategy, stCeloAmount, false);
+        _updateGroupStCelo(strategy, stCeloAmount, false);
     }
 }
diff --git a/test-ts/default-strategy.test.ts b/test-ts/default-strategy.test.ts
@@ -1610,4 +1610,29 @@ describe("DefaultStrategy", () => {
       );
     });
   });
+
+  describe("#updateGroupStCelo", () => {
+    it("should revert when not owner", async () => {
+      await expect(
+        defaultStrategyContract.updateGroupStCelo(groupAddresses[0], 100, true)
+      ).revertedWith("Ownable: caller is not the owner");
+    });
+
+    it("should add group stCelo", async () => {
+      await defaultStrategyContract.connect(owner).updateGroupStCelo(groupAddresses[0], 100, true);
+      const stCelo = await defaultStrategyContract.stCeloInGroup(groupAddresses[0]);
+      expect(stCelo).to.eq(100);
+      const totalStCelo = await defaultStrategyContract.totalStCeloInStrategy();
+      expect(totalStCelo).to.eq(100);
+    });
+
+    it("should subtract group stCelo", async () => {
+      await defaultStrategyContract.connect(owner).updateGroupStCelo(groupAddresses[0], 100, true);
+      await defaultStrategyContract.connect(owner).updateGroupStCelo(groupAddresses[0], 50, false);
+      const stCelo = await defaultStrategyContract.stCeloInGroup(groupAddresses[0]);
+      expect(stCelo).to.eq(50);
+      const totalStCelo = await defaultStrategyContract.totalStCeloInStrategy();
+      expect(totalStCelo).to.eq(50);
+    });
+  });
 });
diff --git a/test-ts/specific_group_strategy.test.ts b/test-ts/specific_group_strategy.test.ts
@@ -1244,4 +1244,35 @@ describe("SpecificGroupStrategy", () => {
       ).revertedWith("Paused()");
     });
   });
+
+  describe("#updateGroupStCelo", () => {
+    it("should revert when not owner", async () => {
+      await expect(
+        specificGroupStrategyContract.updateGroupStCelo(groupAddresses[0], 100, true)
+      ).revertedWith("Ownable: caller is not the owner");
+    });
+
+    it("should add group stCelo", async () => {
+      await specificGroupStrategyContract
+        .connect(owner)
+        .updateGroupStCelo(groupAddresses[0], 100, true);
+      const stCelo = await specificGroupStrategyContract.stCeloInGroup(groupAddresses[0]);
+      expect(stCelo).to.eq(100);
+      const totalStCelo = await specificGroupStrategyContract.totalStCeloLocked();
+      expect(totalStCelo).to.eq(100);
+    });
+
+    it("should subtract group stCelo", async () => {
+      await specificGroupStrategyContract
+        .connect(owner)
+        .updateGroupStCelo(groupAddresses[0], 100, true);
+      await specificGroupStrategyContract
+        .connect(owner)
+        .updateGroupStCelo(groupAddresses[0], 50, false);
+      const stCelo = await specificGroupStrategyContract.stCeloInGroup(groupAddresses[0]);
+      expect(stCelo).to.eq(50);
+      const totalStCelo = await specificGroupStrategyContract.totalStCeloLocked();
+      expect(totalStCelo).to.eq(50);
+    });
+  });
 });