fix: ITS tests

axelarnetwork · Oct 6, 2023 · a2106e3 · a2106e3
1 parent 5a89f3d
commit a2106e3
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 8 deletions.
diff --git a/contracts/interchain-token-service/InterchainTokenService.sol b/contracts/interchain-token-service/InterchainTokenService.sol
@@ -49,6 +49,7 @@ contract InterchainTokenService is
 
     address internal immutable implementationLockUnlock;
     address internal immutable implementationMintBurn;
+    address internal immutable implementationMintBurnFrom;
     address internal immutable implementationLockUnlockFee;
     address internal immutable implementationLiquidityPool;
     IAxelarGasService public immutable gasService;
@@ -98,8 +99,9 @@ contract InterchainTokenService is
 
         if (tokenManagerImplementations.length != uint256(type(TokenManagerType).max) + 1) revert LengthMismatch();
 
-        implementationLockUnlock = _sanitizeTokenManagerImplementation(tokenManagerImplementations, TokenManagerType.LOCK_UNLOCK);
         implementationMintBurn = _sanitizeTokenManagerImplementation(tokenManagerImplementations, TokenManagerType.MINT_BURN);
+        implementationMintBurnFrom = _sanitizeTokenManagerImplementation(tokenManagerImplementations, TokenManagerType.MINT_BURN_FROM);
+        implementationLockUnlock = _sanitizeTokenManagerImplementation(tokenManagerImplementations, TokenManagerType.LOCK_UNLOCK);
         implementationLockUnlockFee = _sanitizeTokenManagerImplementation(
             tokenManagerImplementations,
             TokenManagerType.LOCK_UNLOCK_FEE_ON_TRANSFER
@@ -213,10 +215,12 @@ contract InterchainTokenService is
      */
     function getImplementation(uint256 tokenManagerType) external view returns (address tokenManagerAddress) {
         if (tokenManagerType > uint256(type(TokenManagerType).max)) revert InvalidImplementation();
-        if (TokenManagerType(tokenManagerType) == TokenManagerType.LOCK_UNLOCK) {
-            return implementationLockUnlock;
-        } else if (TokenManagerType(tokenManagerType) == TokenManagerType.MINT_BURN) {
+        if (TokenManagerType(tokenManagerType) == TokenManagerType.MINT_BURN) {
             return implementationMintBurn;
+        } else if (TokenManagerType(tokenManagerType) == TokenManagerType.MINT_BURN_FROM) {
+            return implementationMintBurnFrom;
+        } else if (TokenManagerType(tokenManagerType) == TokenManagerType.LOCK_UNLOCK) {
+            return implementationLockUnlock;
         } else if (TokenManagerType(tokenManagerType) == TokenManagerType.LOCK_UNLOCK_FEE_ON_TRANSFER) {
             return implementationLockUnlockFee;
         } else if (TokenManagerType(tokenManagerType) == TokenManagerType.LIQUIDITY_POOL) {

diff --git a/test/tokenService.js b/test/tokenService.js
@@ -275,6 +275,7 @@ describe('Interchain Token Service', () => {
 
         it('Should revert on invalid token manager implementation length', async () => {
             tokenManagerImplementations.push(wallet);
+
             await expectRevert(
                 (gasOptions) =>
                     deployInterchainTokenService(
@@ -292,6 +293,7 @@ describe('Interchain Token Service', () => {
                 service,
                 'LengthMismatch',
             );
+
             tokenManagerImplementations.pop();
         });
 
@@ -308,18 +310,20 @@ describe('Interchain Token Service', () => {
                 deploymentKey,
             );
 
+            const length = tokenManagerImplementations.length;
             let implementation;
 
-            for (let i = 0; i < 4; i++) {
+            for (let i = 0; i < length; i++) {
                 implementation = await service.getImplementation(i);
                 expect(implementation).to.eq(tokenManagerImplementations[i].address);
             }
 
-            await expectRevert((gasOptions) => service.getImplementation(4, gasOptions), service, 'InvalidImplementation');
+            await expectRevert((gasOptions) => service.getImplementation(length, gasOptions), service, 'InvalidImplementation');
         });
 
         it('Should revert on invalid token manager implementation', async () => {
-            tokenManagerImplementations.pop();
+            const toRemove = tokenManagerImplementations.pop();
+
             await expectRevert(
                 (gasOptions) =>
                     deployInterchainTokenService(
@@ -337,10 +341,14 @@ describe('Interchain Token Service', () => {
                 service,
                 'ZeroAddress',
             );
+
+            tokenManagerImplementations.push(toRemove);
         });
 
         it('Should revert on duplicate token manager type', async () => {
-            tokenManagerImplementations[3] = tokenManagerImplementations[2];
+            const length = tokenManagerImplementations.length;
+            tokenManagerImplementations[length - 1] = tokenManagerImplementations[length - 2];
+
             await expectRevert(
                 (gasOptions) =>
                     deployInterchainTokenService(