From 6454f040e8f0d6d7c4b94feee2be094095c5bd7a Mon Sep 17 00:00:00 2001
From: The3D <frangellaemilio@gmail.com>
Date: Wed, 9 Sep 2020 10:44:34 +0200
Subject: [PATCH] Fixed withdrawal tests

---
 contracts/tokenization/AToken.sol  |  4 ----
 test/helpers/actions.ts            | 11 +++++++----
 test/helpers/utils/calculations.ts | 11 ++++++++---
 test/scenario.spec.ts              |  2 +-
 4 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/contracts/tokenization/AToken.sol b/contracts/tokenization/AToken.sol
index 0b0d7473..ad7d1b07 100644
--- a/contracts/tokenization/AToken.sol
+++ b/contracts/tokenization/AToken.sol
@@ -444,10 +444,6 @@ contract AToken is VersionedInitializable, ERC20, IAToken {
 
     uint256 scaledAmount = amount.rayDiv(index);
 
-    console.log("scaled balanceOf from: %s", scaledBalanceOf(from));
-    console.log("scaled balanceOf to: %s", scaledBalanceOf(to));
-    console.log("scaled amount: %s", scaledAmount);
-
     super._transfer(from, to, scaledAmount);
 
     //if the sender is redirecting his interest towards someone else,
diff --git a/test/helpers/actions.ts b/test/helpers/actions.ts
index be0daac5..e5d6efa6 100644
--- a/test/helpers/actions.ts
+++ b/test/helpers/actions.ts
@@ -56,16 +56,19 @@ const almostEqualOrEqual = function (
       return;
     }
 
+
     this.assert(actual[key] != undefined, `Property ${key} is undefined in the actual data`);
     expect(expected[key] != undefined, `Property ${key} is undefined in the expected data`);
 
+    if (!expected[key] || !actual[key]) {
+      console.log('Found a undefined value for Key ', key, ' value ', expected[key], actual[key]);
+    }
+
     if (actual[key] instanceof BigNumber) {
-      if (!expected[key]) {
-        console.log('Key ', key, ' value ', expected[key], actual[key]);
-      }
+
       const actualValue = (<BigNumber>actual[key]).decimalPlaces(0, BigNumber.ROUND_DOWN);
       const expectedValue = (<BigNumber>expected[key]).decimalPlaces(0, BigNumber.ROUND_DOWN);
-
+  
       this.assert(
         actualValue.eq(expectedValue) ||
           actualValue.plus(1).eq(expectedValue) ||
diff --git a/test/helpers/utils/calculations.ts b/test/helpers/utils/calculations.ts
index f66117d9..fcdece3e 100644
--- a/test/helpers/utils/calculations.ts
+++ b/test/helpers/utils/calculations.ts
@@ -120,7 +120,9 @@ export const calcExpectedUserDataAfterWithdraw = (
   }
 
   expectedUserData.scaledATokenBalance = calcExpectedScaledATokenBalance(reserveDataAfterAction, userDataBeforeAction, new BigNumber(0), new BigNumber(amountWithdrawn));
-  
+
+  console.log("Scaled balance is ", expectedUserData.scaledATokenBalance.toFixed());
+
   expectedUserData.currentATokenBalance = aTokenBalance.minus(
     amountWithdrawn
   );
@@ -161,8 +163,10 @@ export const calcExpectedUserDataAfterWithdraw = (
 
   if (expectedUserData.currentATokenBalance.eq(0) && expectedUserData.redirectedBalance.eq(0)) {
     expectedUserData.interestRedirectionAddress = ZERO_ADDRESS;
+    expectedUserData.interestRedirectionIndex = new BigNumber(0);
   } else {
     expectedUserData.interestRedirectionAddress = userDataBeforeAction.interestRedirectionAddress;
+    expectedUserData.interestRedirectionIndex = userDataBeforeAction.interestRedirectionAddress == ZERO_ADDRESS ? new BigNumber(0) : reserveDataAfterAction.liquidityIndex;
   }
 
   expectedUserData.redirectionAddressRedirectedBalance = calcExpectedRedirectedBalance(
@@ -172,7 +176,7 @@ export const calcExpectedUserDataAfterWithdraw = (
     new BigNumber(0),
     new BigNumber(amountWithdrawn)
   );
-
+  
   return expectedUserData;
 };
 
@@ -564,9 +568,10 @@ export const calcExpectedUserDataAfterBorrow = (
     userDataBeforeAction,
     currentTimestamp
   );
-  expectedUserData.principalATokenBalance = userDataBeforeAction.principalATokenBalance;
+  expectedUserData.scaledATokenBalance = userDataBeforeAction.scaledATokenBalance;
   expectedUserData.redirectedBalance = userDataBeforeAction.redirectedBalance;
   expectedUserData.interestRedirectionAddress = userDataBeforeAction.interestRedirectionAddress;
+  expectedUserData.interestRedirectionIndex = userDataBeforeAction.interestRedirectionIndex;
   expectedUserData.redirectionAddressRedirectedBalance =
     userDataBeforeAction.redirectionAddressRedirectedBalance;
   expectedUserData.currentATokenUserIndex = userDataBeforeAction.currentATokenUserIndex;
diff --git a/test/scenario.spec.ts b/test/scenario.spec.ts
index 471bb9ab..5d449d76 100644
--- a/test/scenario.spec.ts
+++ b/test/scenario.spec.ts
@@ -12,7 +12,7 @@ BigNumber.config({DECIMAL_PLACES: 0, ROUNDING_MODE: BigNumber.ROUND_DOWN});
 
 const scenarioFolder = './test/helpers/scenarios/';
 
-const selectedScenarios: string[] = ['withdraw.json'];
+const selectedScenarios: string[] = [];
 
 fs.readdirSync(scenarioFolder).forEach((file) => {
   if (selectedScenarios.length > 0 && !selectedScenarios.includes(file)) return;