fix: ability to upgrade user token

VenusProtocol · Sep 20, 2023 · a19e29c · a19e29c
1 parent f7d463b
commit a19e29c
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 2 deletions.
diff --git a/contracts/Tokens/Prime/Prime.sol b/contracts/Tokens/Prime/Prime.sol
@@ -88,6 +88,9 @@ contract Prime is IIncomeDestination, AccessControlledV8, PausableUpgradeable, M
     /// @notice Emitted when interest is claimed
     event InterestClaimed(address indexed user, address indexed market, uint256 amount);
 
+    /// @notice Emitted when revocable token is upgraded to irrevocable token
+    event TokenUpgraded(address indexed user);
+
     /// @custom:oz-upgrades-unsafe-allow constructor
     /**
      * @notice Prime constructor
@@ -265,8 +268,18 @@ contract Prime is IIncomeDestination, AccessControlledV8, PausableUpgradeable, M
 
         if (isIrrevocable) {
             for (uint i = 0; i < users.length; ) {
-                _mint(true, users[i]);
-                _initializeMarkets(users[i]);
+                Token storage userToken = tokens[users[i]];
+                if (userToken.exists && !userToken.isIrrevocable) {
+                    //upgrade to irrevocable token
+                    userToken.isIrrevocable = true;      
+                    totalIrrevocable++;
+                    totalRevocable--; 
+
+                    emit TokenUpgraded(users[i]);
+                } else {
+                    _mint(true, users[i]);
+                    _initializeMarkets(users[i]);
+                }
 
                 unchecked {
                     i++;

diff --git a/tests/hardhat/Prime/Prime.ts b/tests/hardhat/Prime/Prime.ts
@@ -418,6 +418,28 @@ describe("PrimeScenario Token", () => {
       expect(token.isIrrevocable).to.be.equal(false);
       expect(token.exists).to.be.equal(true);
     });
+
+    it("upgrade", async () => {
+      await prime.issue(false, [user1.getAddress(), user2.getAddress()]);
+
+      let token = await prime.tokens(user1.getAddress());
+      expect(token.exists).to.be.equal(true);
+      expect(token.isIrrevocable).to.be.equal(false);
+
+      token = await prime.tokens(user2.getAddress());
+      expect(token.isIrrevocable).to.be.equal(false);
+      expect(token.exists).to.be.equal(true);
+
+      await prime.issue(true, [user1.getAddress()]);
+
+      token = await prime.tokens(user1.getAddress());
+      expect(token.isIrrevocable).to.be.equal(true);
+      expect(token.exists).to.be.equal(true);
+
+      token = await prime.tokens(user2.getAddress());
+      expect(token.isIrrevocable).to.be.equal(false);
+      expect(token.exists).to.be.equal(true);
+    });
   });
 
   describe("boosted yield", () => {