diff --git a/contracts/token-manager/implementations/TokenManagerLockUnlockFeeOnTransfer.sol b/contracts/token-manager/implementations/TokenManagerLockUnlockFee.sol
similarity index 89%
rename from contracts/token-manager/implementations/TokenManagerLockUnlockFeeOnTransfer.sol
rename to contracts/token-manager/implementations/TokenManagerLockUnlockFee.sol
index 733ec6aa..b3b95f01 100644
--- a/contracts/token-manager/implementations/TokenManagerLockUnlockFeeOnTransfer.sol
+++ b/contracts/token-manager/implementations/TokenManagerLockUnlockFee.sol
@@ -47,15 +47,11 @@ contract TokenManagerLockUnlockFee is TokenManager, NoReEntrancy {
      */
     function _takeToken(address from, uint256 amount) internal override noReEntrancy returns (uint256) {
         IERC20 token = IERC20(tokenAddress());
-        uint256 balance = token.balanceOf(address(this));
+        uint256 balanceBefore = token.balanceOf(address(this));
 
         token.safeTransferFrom(from, address(this), amount);
 
-        uint256 diff = token.balanceOf(address(this)) - balance;
-        if (diff < amount) {
-            amount = diff;
-        }
-        return amount;
+        return token.balanceOf(address(this)) - balanceBefore;
     }
 
     /**
@@ -66,15 +62,11 @@ contract TokenManagerLockUnlockFee is TokenManager, NoReEntrancy {
      */
     function _giveToken(address to, uint256 amount) internal override noReEntrancy returns (uint256) {
         IERC20 token = IERC20(tokenAddress());
-        uint256 balance = token.balanceOf(to);
+        uint256 balanceBefore = token.balanceOf(to);
 
         token.safeTransfer(to, amount);
 
-        uint256 diff = token.balanceOf(to) - balance;
-        if (diff < amount) {
-            amount = diff;
-        }
-        return amount;
+        return token.balanceOf(to) - balanceBefore;
     }
 
     /**