Merge ff79541 into 89535da

pooltogether · Jul 6, 2021 · 7847957 · 7847957
2 parents 89535da + ff79541
commit 7847957
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 19 deletions.
diff --git a/contracts/yield-source/ATokenYieldSource.sol b/contracts/yield-source/ATokenYieldSource.sol
@@ -203,15 +203,18 @@ contract ATokenYieldSource is ERC20Upgradeable, IProtocolYieldSource, AssetManag
   /// @param redeemAmount The amount of asset tokens to be redeemed
   /// @return The actual amount of asset tokens that were redeemed
   function redeemToken(uint256 redeemAmount) external override nonReentrant returns (uint256) {
+    address _tokenAddress = _tokenAddress();
+    IERC20Upgradeable _depositToken = IERC20Upgradeable(_tokenAddress);
+
     uint256 shares = _tokenToShares(redeemAmount);
     _burn(msg.sender, shares);
 
-    uint256 beforeBalance = aToken.balanceOf(address(this));
-    _lendingPool().withdraw(address(_tokenAddress()), redeemAmount, address(this));
-    uint256 afterBalance = aToken.balanceOf(address(this));
+    uint256 beforeBalance = _depositToken.balanceOf(address(this));
+    _lendingPool().withdraw(_tokenAddress, redeemAmount, address(this));
+    uint256 afterBalance = _depositToken.balanceOf(address(this));
 
-    uint256 balanceDiff = beforeBalance.sub(afterBalance);
-    IERC20Upgradeable(_tokenAddress()).safeTransfer(msg.sender, balanceDiff);
+    uint256 balanceDiff = afterBalance.sub(beforeBalance);
+    _depositToken.safeTransfer(msg.sender, balanceDiff);
 
     emit RedeemedToken(msg.sender, shares, redeemAmount);
     return balanceDiff;

diff --git a/test/ATokenYieldSource.test.ts b/test/ATokenYieldSource.test.ts
@@ -26,7 +26,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;
@@ -93,12 +93,12 @@ describe('ATokenYieldSource', () => {
     const hardhatATokenYieldSourceHarness = await ATokenYieldSource.deploy()
 
     aTokenYieldSource = (await ethers.getContractAt(
-        'ATokenYieldSourceHarness',
-        hardhatATokenYieldSourceHarness.address,
-        contractsOwner,
-      )) as ATokenYieldSourceHarness;
-  
-    const initializeTx = await aTokenYieldSource.initialize(
+      'ATokenYieldSourceHarness',
+      hardhatATokenYieldSourceHarness.address,
+      contractsOwner
+    )) as unknown as ATokenYieldSourceHarness;
+
+    await aTokenYieldSource.initialize(
       aToken.address,
       lendingPoolAddressesProviderRegistry.address,
       18,
@@ -173,8 +173,7 @@ describe('ATokenYieldSource', () => {
       expect(await aTokenYieldSource.tokenToShares(toWei('1'))).to.equal(2);
     });
 
-    it('should fail to return shares if aToken total supply increases too much', async () => { // failing here
-
+    it('should fail to return shares if aToken total supply increases too much', async () => {
       await aTokenYieldSource.mint(yieldSourceOwner.address, toWei('100'));
       await aTokenYieldSource.mint(wallet2.address, toWei('100'));
       await aToken.mock.balanceOf.withArgs(aTokenYieldSource.address).returns(toWei('100'));
@@ -183,7 +182,6 @@ describe('ATokenYieldSource', () => {
 
       await aToken.mock.balanceOf.withArgs(aTokenYieldSource.address).returns(ethers.utils.parseUnits('100', 37));
       await expect(aTokenYieldSource.supplyTokenTo(toWei('1'), wallet2.address)).to.be.revertedWith('ATokenYieldSource/shares-equal-zero');
-
     });
   });
 
@@ -287,15 +285,16 @@ describe('ATokenYieldSource', () => {
       await aToken.mock.balanceOf
         .withArgs(aTokenYieldSource.address)
         .returns(yieldSourceOwnerBalance);
+      await underlyingToken.mock.balanceOf
+        .withArgs(aTokenYieldSource.address)
+        .returns(redeemAmount);
       await lendingPool.mock.withdraw
         .withArgs(underlyingToken.address, redeemAmount, aTokenYieldSource.address)
         .returns(redeemAmount);
 
-      const balanceAfter = await aToken.balanceOf(aTokenYieldSource.address);
-      const balanceDiff = yieldSourceOwnerBalance.sub(balanceAfter);
-
+      // mocked underlyingToken will always return 100 so afterBalance.sub(beforeBalance) returns 0
       await underlyingToken.mock.transfer
-        .withArgs(yieldSourceOwner.address, balanceDiff)
+        .withArgs(yieldSourceOwner.address, toWei('0'))
         .returns(true);
 
       await aTokenYieldSource.connect(yieldSourceOwner).redeemToken(redeemAmount);