returning fee amount

The-Poolz · Mar 25, 2024 · eedd013 · eedd013
1 parent ca893c4
commit eedd013
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 6 deletions.
diff --git a/contracts/Fee/FeeBaseHelper.sol b/contracts/Fee/FeeBaseHelper.sol
@@ -18,16 +18,16 @@ abstract contract FeeBaseHelper is ERC20Helper, WhiteListHelper {
 
     mapping(address => uint) public FeeReserve;
 
-    function TakeFee() internal virtual firewallProtected {
-        uint feeToPay = FeeAmount;
-        if(feeToPay == 0) return;
+    function TakeFee() internal virtual firewallProtected returns(uint feeToPay){
+        feeToPay = FeeAmount;
+        if(feeToPay == 0) return 0;
         uint credits = getCredits(msg.sender);
         if(credits > 0) {
             _whiteListRegister(msg.sender, credits < feeToPay ? credits : feeToPay);
             if(credits < feeToPay) {
                 feeToPay -= credits;
             } else {
-                return;
+                return 0;
             }
         }
         _TakeFee(feeToPay);

diff --git a/contracts/mocks/FeeHelper.sol b/contracts/mocks/FeeHelper.sol
@@ -5,7 +5,7 @@ import "../Fee/FeeBaseHelper.sol";
 
 contract FeeHelper is FeeBaseHelper {
 
-    function MethodWithFee() public payable {
-        TakeFee();
+    function MethodWithFee() public payable returns(uint){
+        return TakeFee();
     }
 }
diff --git a/test/3_Fee.ts b/test/3_Fee.ts
@@ -24,12 +24,14 @@ describe('Fee Helper Test', function () {
   it('zero fee', async () => {
     const gasPrice = 2500000000;
     const oldBal = await payer.getBalance();
+    const feePaid = await feeHelper.connect(payer).callStatic.MethodWithFee({ gasPrice: gasPrice });
     const tx = await feeHelper.connect(payer).MethodWithFee({ gasPrice: gasPrice });
     const txReceipt = await tx.wait();
     const actualBal = await payer.getBalance();
     const gas = txReceipt.gasUsed.mul(gasPrice);
     const expectedBal = oldBal.sub(gas);
     expect(actualBal).to.be.equal(expectedBal);
+    expect(feePaid).to.be.equal(0);
   });
 
   describe('test ERC20 token', async () => {
@@ -44,9 +46,11 @@ describe('Fee Helper Test', function () {
     it('should pay', async () => {
       await token.connect(payer).approve(feeHelper.address, fee);
       const oldBal = await token.balanceOf(feeHelper.address);
+      const feePaid = await feeHelper.connect(payer).callStatic.MethodWithFee();
       await feeHelper.connect(payer).MethodWithFee();
       const actualBal = await token.balanceOf(feeHelper.address);
       expect(actualBal).to.be.equal(oldBal.add(fee));
+      expect(feePaid).to.be.equal(fee);
     });
 
     it('should withdraw', async () => {
@@ -67,9 +71,11 @@ describe('Fee Helper Test', function () {
 
     it('should pay', async () => {
       const oldBal = await ethers.provider.getBalance(feeHelper.address);
+      const feePaid = await feeHelper.connect(payer).callStatic.MethodWithFee({ value: fee });
       await feeHelper.connect(payer).MethodWithFee({ value: fee });
       const actualBal = await ethers.provider.getBalance(feeHelper.address);
       expect(actualBal).to.be.equal(oldBal.add(fee));
+      expect(feePaid).to.be.equal(fee);
     });
 
     it("should revert when fee is not paid", async () => {
@@ -116,10 +122,12 @@ describe('Fee Helper Test', function () {
       await feeHelper.addUsers([payer.address], [0])
       const beforePayerBalance = await token.balanceOf(payer.address);
       await token.connect(payer).approve(feeHelper.address, fee);
+      const feePaid = await feeHelper.connect(payer).callStatic.MethodWithFee();
       const tx = await feeHelper.connect(payer).MethodWithFee();
       const afterPayerBalance = await token.balanceOf(payer.address);
       expect(beforePayerBalance.sub(afterPayerBalance)).to.be.equal(fee);
       expect(await feeHelper.FeeToken()).to.be.equal(token.address);
+      expect(feePaid).to.be.equal(fee);
       await expect(tx).to.emit(feeHelper, "TransferIn").withArgs(fee, payer.address, token.address);
     })
 
@@ -128,10 +136,12 @@ describe('Fee Helper Test', function () {
       const credits = fee * 5;
       await feeHelper.addUsers([payer.address], [credits])
       const beforePayerBalance = await token.balanceOf(payer.address);
+      const feePaid = await feeHelper.connect(payer).callStatic.MethodWithFee();
       const tx = await feeHelper.connect(payer).MethodWithFee();
       const afterPayerBalance = await token.balanceOf(payer.address);
       expect(beforePayerBalance).to.be.equal(afterPayerBalance);
       expect(await feeHelper.FeeToken()).to.be.equal(token.address);
+      expect(feePaid).to.be.equal(0);
       await expect(tx).to.not.emit(feeHelper, "TransferIn");
     })
 
@@ -141,10 +151,12 @@ describe('Fee Helper Test', function () {
       await feeHelper.addUsers([payer.address], [credits])
       const beforePayerBalance = await token.balanceOf(payer.address);
       await token.connect(payer).approve(feeHelper.address, fee - credits);
+      const feePaid = await feeHelper.connect(payer).callStatic.MethodWithFee();
       const tx = await feeHelper.connect(payer).MethodWithFee();
       const afterPayerBalance = await token.balanceOf(payer.address);
       expect(beforePayerBalance.sub(afterPayerBalance)).to.be.equal(fee - credits);
       expect(await feeHelper.FeeToken()).to.be.equal(token.address);
+      expect(feePaid).to.be.equal(fee - credits);
       await expect(tx).to.emit(feeHelper, "TransferIn").withArgs(fee - credits, payer.address, token.address);
     })