From e00fd1381616c2373acf74e2f30aec467d1b7468 Mon Sep 17 00:00:00 2001
From: Lukas <lukasstrassel@googlemail.com>
Date: Wed, 21 Sep 2022 17:48:29 +0200
Subject: [PATCH] feat: add repay method (#14)

* feat: add repay method

* tests: add more tests

* fix: update tests to use msg.sender
---
 src/ProtocolV3TestBase.sol | 58 ++++++++++++++++++++++++++++++++------
 1 file changed, 50 insertions(+), 8 deletions(-)

diff --git a/src/ProtocolV3TestBase.sol b/src/ProtocolV3TestBase.sol
index 2275a87f0..7844858d4 100644
--- a/src/ProtocolV3TestBase.sol
+++ b/src/ProtocolV3TestBase.sol
@@ -83,7 +83,7 @@ contract ProtocolV3TestBase is Test {
 
   function e2eTest(IPool pool) public {
     ReserveConfig[] memory configs = _getReservesConfigs(pool);
-    deal(address(this), 1000 ether);
+    deal(msg.sender, 1000 ether);
     uint256 snapshot = vm.snapshot();
     _supplyWithdrawFlow(configs, pool);
     vm.revertTo(snapshot);
@@ -146,7 +146,7 @@ contract ProtocolV3TestBase is Test {
     for (uint256 i = 0; i < configs.length; i++) {
       uint256 amount = 10**configs[i].decimals;
       if (configs[i].borrowingEnabled) {
-        _borrow(configs[i], pool, amount, false);
+        this._borrow(configs[i], pool, amount, false);
       } else {
         console.log('SKIP: BORROWING_DISABLED %s', configs[i].symbol);
       }
@@ -163,7 +163,7 @@ contract ProtocolV3TestBase is Test {
     for (uint256 i = 0; i < configs.length; i++) {
       uint256 amount = 10**configs[i].decimals;
       if (configs[i].borrowingEnabled && configs[i].stableBorrowRateEnabled) {
-        _borrow(configs[i], pool, amount, true);
+        this._borrow(configs[i], pool, amount, true);
       } else {
         console.log('SKIP: STABLE_BORROWING_DISABLED %s', configs[i].symbol);
       }
@@ -175,10 +175,13 @@ contract ProtocolV3TestBase is Test {
     IPool pool,
     uint256 amount
   ) internal {
-    deal(config.underlying, address(this), amount);
+    uint256 aTokenBefore = IERC20(config.aToken).balanceOf(msg.sender);
+    deal(config.underlying, msg.sender, amount);
     IERC20(config.underlying).approve(address(pool), amount);
     console.log('SUPPLY: %s, Amount: %s', config.symbol, amount);
-    pool.deposit(config.underlying, amount, address(this), 0);
+    pool.deposit(config.underlying, amount, msg.sender, 0);
+    uint256 aTokenAfter = IERC20(config.aToken).balanceOf(msg.sender);
+    require(_almostEqual(aTokenAfter, aTokenBefore + amount), '_deposit() : ERROR');
   }
 
   function _withdraw(
@@ -187,12 +190,22 @@ contract ProtocolV3TestBase is Test {
     uint256 amount,
     bool max
   ) internal returns (uint256) {
+    uint256 aTokenBefore = IERC20(config.aToken).balanceOf(msg.sender);
     uint256 amountOut = pool.withdraw(
       config.underlying,
       max ? type(uint256).max : amount,
-      address(this)
+      msg.sender
     );
     console.log('WITHDRAW: %s, Amount: %s', config.symbol, amountOut);
+    uint256 aTokenAfter = IERC20(config.aToken).balanceOf(msg.sender);
+    if (aTokenBefore < amount || max) {
+      require(aTokenAfter == 0, '_widthdraw(): DUST_AFTER_WITHDRAW_ALL');
+    } else {
+      require(
+        _almostEqual(aTokenAfter, aTokenBefore - amount),
+        '_withdraw() : INCONSISTENT_ATOKEN_BALANCE_AFTER'
+      );
+    }
     return amountOut;
   }
 
@@ -201,9 +214,29 @@ contract ProtocolV3TestBase is Test {
     IPool pool,
     uint256 amount,
     bool stable
-  ) internal {
+  ) external {
+    address debtToken = stable ? config.stableDebtToken : config.variableDebtToken;
+    uint256 debtBefore = IERC20(debtToken).balanceOf(msg.sender);
     console.log('BORROW: %s, Amount %s, Stable: %s', config.symbol, amount, stable);
-    pool.borrow(config.underlying, amount, stable ? 1 : 2, 0, address(this));
+    pool.borrow(config.underlying, amount, stable ? 1 : 2, 0, msg.sender);
+    uint256 debtAfter = IERC20(debtToken).balanceOf(msg.sender);
+    require(_almostEqual(debtAfter, debtBefore + amount), '_borrow() : ERROR');
+  }
+
+  function _repay(
+    ReserveConfig memory config,
+    IPool pool,
+    uint256 amount,
+    bool stable
+  ) internal {
+    address debtToken = stable ? config.stableDebtToken : config.variableDebtToken;
+    uint256 debtBefore = IERC20(debtToken).balanceOf(msg.sender);
+    deal(config.underlying, msg.sender, amount);
+    IERC20(config.underlying).approve(address(pool), amount);
+    console.log('REPAY: %s, Amount: %s', config.symbol, amount);
+    pool.repay(config.underlying, amount, stable ? 1 : 2, msg.sender);
+    uint256 debtAfter = IERC20(debtToken).balanceOf(msg.sender);
+    require(debtAfter == ((debtBefore > amount) ? debtBefore - amount : 0), '_repay() : ERROR');
   }
 
   function _isInUint256Array(uint256[] memory haystack, uint256 needle)
@@ -856,4 +889,13 @@ contract ProtocolV3TestBase is Test {
       revert('_getAssetOnEmodeCategory(): LESS_ASSETS_IN_CATEGORY_THAN_EXPECTED');
     }
   }
+
+  /// @dev To contemplate +1/-1 precision issues when rounding, mainly on aTokens
+  function _almostEqual(uint256 a, uint256 b) internal pure returns (bool) {
+    if (b == 0) {
+      return (a == b) || (a == (b + 1));
+    } else {
+      return (a == b) || (a == (b + 1)) || (a == (b - 1));
+    }
+  }
 }