Added some tests and fixed constants

axelarnetwork · Oct 19, 2023 · fc0d72b · fc0d72b
1 parent afdb98d
commit fc0d72b
Show file tree

Hide file tree

Showing 4 changed files with 76 additions and 6 deletions.
diff --git a/contracts/interfaces/ITokenManager.sol b/contracts/interfaces/ITokenManager.sol
@@ -18,7 +18,8 @@ interface ITokenManager is ITokenManagerType, IOperatable, IFlowLimit, IImplemen
     error GiveTokenFailed();
     error NotToken();
     error ZeroAddress();
-    error NotFlowLimiter();
+    error AlreadyFlowLimiter(address flowLimiter);
+    error NotFlowLimiter(address flowLimiter);
 
     /**
      * @notice A function that returns the token id.

diff --git a/contracts/token-manager/TokenManager.sol b/contracts/token-manager/TokenManager.sol
@@ -205,6 +205,7 @@ abstract contract TokenManager is ITokenManager, Operatable, FlowLimit, Implemen
      */
     function addFlowLimiter(address flowLimiter) external onlyRole(OPERATOR) {
         if (flowLimiter == address(0)) revert ZeroAddress();
+        if (hasRole(flowLimiter, FLOW_LIMITER)) revert AlreadyFlowLimiter(flowLimiter);
         uint8[] memory roles = new uint8[](1);
         roles[0] = FLOW_LIMITER;
         _addRoles(flowLimiter, roles);
@@ -216,7 +217,7 @@ abstract contract TokenManager is ITokenManager, Operatable, FlowLimit, Implemen
      */
     function removeFlowLimiter(address flowLimiter) external onlyRole(OPERATOR) {
         if (flowLimiter == address(0)) revert ZeroAddress();
-        if (!hasRole(flowLimiter, FLOW_LIMITER)) revert NotFlowLimiter();
+        if (!hasRole(flowLimiter, FLOW_LIMITER)) revert NotFlowLimiter(flowLimiter);
         uint8[] memory roles = new uint8[](1);
         roles[0] = FLOW_LIMITER;
         _removeRoles(flowLimiter, roles);

diff --git a/contracts/utils/RolesConstants.sol b/contracts/utils/RolesConstants.sol
@@ -4,6 +4,6 @@ pragma solidity ^0.8.0;
 
 contract RolesConstants {
     uint8 internal constant DISTRIBUTOR = 1;
-    uint8 internal constant OPERATOR = 1;
-    uint8 internal constant FLOW_LIMITER = 1;
+    uint8 internal constant OPERATOR = 2;
+    uint8 internal constant FLOW_LIMITER = 3;
 }
diff --git a/test/TokenService.js b/test/TokenService.js
@@ -33,7 +33,9 @@ const MINT_BURN_FROM = 1;
 const LOCK_UNLOCK = 2;
 const LOCK_UNLOCK_FEE_ON_TRANSFER = 3;
 
-const OPERATOR_ROLE = 1;
+const DISTRIBUTOR_ROLE = 1;
+const OPERATOR_ROLE = 2;
+const FLOW_LIMITER_ROLE = 3;
 
 describe('Interchain Token Service', () => {
     let wallet, otherWallet;
@@ -2030,7 +2032,7 @@ describe('Interchain Token Service', () => {
                 tokenIds.push(tokenId);
                 tokenManagers.push(tokenManager);
 
-                await tokenManager.transferOperatorship(service.address).then((tx) => tx.wait());
+                await tokenManager.addFlowLimiter(service.address).then((tx) => tx.wait());
             }
 
             const flowLimits = new Array(tokenManagers.length).fill(flowLimit);
@@ -2054,4 +2056,70 @@ describe('Interchain Token Service', () => {
             await expectRevert((gasOptions) => service.setFlowLimits(tokenIds, flowLimits, gasOptions), service, 'LengthMismatch');
         });
     });
+
+    describe('Flow Limiters', () => {
+        let tokenManager, tokenId;
+        const sendAmount = 1234;
+        const flowLimit = (sendAmount * 3) / 2;
+        const mintAmount = flowLimit * 3;
+
+        beforeEach(async () => {
+            [, tokenManager, tokenId] = await deployFunctions.mintBurn(`Test Token Lock Unlock`, 'TT', 12, mintAmount);
+        });
+
+        it('Should have only the owner be a flow limiter', async() => {
+            expect(await tokenManager.hasRole(wallet.address, FLOW_LIMITER_ROLE)).to.equal(true);
+            expect(await tokenManager.hasRole(otherWallet.address, FLOW_LIMITER_ROLE)).to.equal(false);
+        });
+
+        it('Should be able to add a flow limiter', async() => {
+            await expect(tokenManager.addFlowLimiter(otherWallet.address))
+                .to.emit(tokenManager, 'RolesAdded')
+                .withArgs(otherWallet.address, [FLOW_LIMITER_ROLE]);
+
+
+            expect(await tokenManager.hasRole(wallet.address, FLOW_LIMITER_ROLE)).to.equal(true);
+            expect(await tokenManager.hasRole(otherWallet.address, FLOW_LIMITER_ROLE)).to.equal(true);
+        });
+
+        it('Should be able to remove a flow limiter', async() => {
+            await expect(tokenManager.removeFlowLimiter(wallet.address))
+                .to.emit(tokenManager, 'RolesRemoved')
+                .withArgs(wallet.address, [FLOW_LIMITER_ROLE]);
+
+
+            expect(await tokenManager.hasRole(wallet.address, FLOW_LIMITER_ROLE)).to.equal(false);
+            expect(await tokenManager.hasRole(otherWallet.address, FLOW_LIMITER_ROLE)).to.equal(false);
+        });
+
+        it('Should revert if trying to add a flow limiter as not the operator', async() => {
+            await expectRevert(
+                (gasOptions) => tokenManager.connect(otherWallet).addFlowLimiter(otherWallet.address, gasOptions),
+                tokenManager, 
+                'MissingRole');
+        });
+
+        it('Should revert if trying to add a flow limiter as not the operator', async() => {
+            await expectRevert(
+                (gasOptions) => tokenManager.connect(otherWallet).removeFlowLimiter(wallet.address, gasOptions),
+                tokenManager, 
+                'MissingRole');
+        });
+
+        it('Should revert if trying to add an existing flow limiter', async() => {
+            await expectRevert(
+                (gasOptions) => tokenManager.addFlowLimiter(wallet.address, gasOptions),
+                tokenManager, 
+                'AlreadyFlowLimiter');
+        });
+
+        it('Should revert if trying to add a flow limiter as not the operator', async() => {
+            await expectRevert(
+                (gasOptions) => tokenManager.removeFlowLimiter(otherWallet.address, gasOptions),
+                tokenManager, 
+                'NotFlowLimiter');
+        });
+
+
+    })
 });