feat(contract): add approveMaxAmount function

pooltogether · Jul 6, 2021 · 5378d9e · 5378d9e
1 parent 4fefb18
commit 5378d9e
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 16 deletions.
diff --git a/contracts/yield-source/ATokenYieldSource.sol b/contracts/yield-source/ATokenYieldSource.sol
@@ -70,8 +70,8 @@ contract ATokenYieldSource is ERC20Upgradeable, IProtocolYieldSource, AssetManag
   /// @notice Interface for Aave lendingPoolAddressesProviderRegistry
   ILendingPoolAddressesProviderRegistry public lendingPoolAddressesProviderRegistry;
 
-  
-  /// @notice Mock Initializer to prevent 
+
+  /// @notice Mock Initializer to prevent
   function freeze() public initializer {
     //no-op
   }
@@ -103,10 +103,10 @@ contract ATokenYieldSource is ERC20Upgradeable, IProtocolYieldSource, AssetManag
     __ERC20_init(_name,_symbol);
     require(_decimals > 0, "ATokenYieldSource/decimals-gt-zero");
     _setupDecimals(_decimals);
-    
+
     // approve once for max amount
     IERC20Upgradeable(_tokenAddress()).safeApprove(address(_lendingPool()), type(uint256).max);
-    
+
     emit ATokenYieldSourceInitialized (
       _aToken,
       _lendingPoolAddressesProviderRegistry,
@@ -119,6 +119,14 @@ contract ATokenYieldSource is ERC20Upgradeable, IProtocolYieldSource, AssetManag
     return true;
   }
 
+  /// @notice Approve lending pool contract to spend max uint256 amount
+  /// @dev Emergency function to re-approve max amount if approval amount dropped too low
+  /// @return true if operation is successful
+  function approveMaxAmount() external returns (bool) {
+    IERC20Upgradeable(_tokenAddress()).safeApprove(address(_lendingPool()), type(uint256).max);
+    return true;
+  }
+
   /// @notice Returns the ERC20 asset token used for deposits
   /// @return The ERC20 asset token address
   function depositToken() public view override returns (address) {
@@ -175,12 +183,11 @@ contract ATokenYieldSource is ERC20Upgradeable, IProtocolYieldSource, AssetManag
 
   /// @notice Deposit asset tokens to Aave
   /// @param mintAmount The amount of asset tokens to be deposited
-  /// @return 0 if successful 
+  /// @return 0 if successful
   function _depositToAave(uint256 mintAmount) internal returns (uint256) {
     IERC20Upgradeable _depositToken = IERC20Upgradeable(_tokenAddress());
 
     _depositToken.safeTransferFrom(msg.sender, address(this), mintAmount);
-
     _lendingPool().deposit(address(_tokenAddress()), mintAmount, address(this), _getRefferalCode());
     return 0;
   }

diff --git a/test/ATokenYieldSource.test.ts b/test/ATokenYieldSource.test.ts
@@ -27,7 +27,7 @@ import SafeERC20WrapperUpgradeable from '../abis/SafeERC20WrapperUpgradeable.jso
 const toWei = ethers.utils.parseEther;
 
 describe('ATokenYieldSource', () => {
-  let contractsOwner: Signer;
+  let contractsOwner: SignerWithAddress;
   let yieldSourceOwner: SignerWithAddress;
   let wallet2: SignerWithAddress;
   let provider: JsonRpcProvider;
@@ -94,16 +94,15 @@ describe('ATokenYieldSource', () => {
     const hardhatATokenYieldSourceHarness = await ATokenYieldSource.deploy()
 
     aTokenYieldSource = (await ethers.getContractAt(
-        'ATokenYieldSourceHarness',
-        hardhatATokenYieldSourceHarness.address,
-        contractsOwner,
-      )) as ATokenYieldSourceHarness;
-    
+      'ATokenYieldSourceHarness',
+      hardhatATokenYieldSourceHarness.address,
+      contractsOwner
+    )) as unknown as ATokenYieldSourceHarness;
+
     await underlyingToken.mock.allowance.returns(ethers.constants.Zero);
     await underlyingToken.mock.approve.withArgs(lendingPool.address, ethers.constants.MaxUint256).returns(true);
-
-
-    const initializeTx = await aTokenYieldSource.initialize(
+
+    await aTokenYieldSource.initialize(
       aToken.address,
       lendingPoolAddressesProviderRegistry.address,
       18,
@@ -124,6 +123,19 @@ describe('ATokenYieldSource', () => {
     });
   });
 
+  describe('depositToken()', () => {
+    it('should approve lending pool to spend max uint256 amount', async () => {
+      await underlyingToken.mock.allowance.returns(ethers.constants.Zero);
+      expect(await underlyingToken.allowance(aTokenYieldSource.address, lendingPool.address)).to.equal(ethers.constants.Zero);
+
+      await underlyingToken.mock.approve.withArgs(lendingPool.address, ethers.constants.MaxUint256).returns(true);
+      expect(await aTokenYieldSource.callStatic.approveMaxAmount()).to.equal(true);
+
+      await underlyingToken.mock.allowance.returns(ethers.constants.MaxUint256);
+      expect(await underlyingToken.allowance(aTokenYieldSource.address, lendingPool.address)).to.equal(ethers.constants.MaxUint256);
+    });
+  });
+
   describe('depositToken()', () => {
     it('should return the underlying token', async () => {
       expect(await aTokenYieldSource.depositToken()).to.equal(underlyingToken.address);
@@ -180,7 +192,7 @@ describe('ATokenYieldSource', () => {
     });
 
     it('should fail to return shares if aToken total supply increases too much', async () => { // failing here
-    
+
       await aTokenYieldSource.mint(yieldSourceOwner.address, toWei('100'));
       await aTokenYieldSource.mint(wallet2.address, toWei('100'));
       await aToken.mock.balanceOf.withArgs(aTokenYieldSource.address).returns(toWei('100'));