check that balances are changed after rToken mint/redeem

omni · Jan 8, 2020 · d11a9fb · d11a9fb
1 parent d1eb3eb
commit d11a9fb
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 0 deletions.
diff --git a/contracts/interfaces/IRToken.sol b/contracts/interfaces/IRToken.sol
@@ -9,4 +9,5 @@ contract IRToken {
     function redeem(uint256 redeemTokens) external returns (bool);
     function createHat(address[] recipients, uint32[] proportions, bool doChangeHat) external returns (uint256 hatID);
     function redeemAll() external returns (bool);
+    function balanceOf(address owner) external view returns (uint256);
 }
diff --git a/contracts/upgradeable_contracts/RTokenConnector.sol b/contracts/upgradeable_contracts/RTokenConnector.sol
@@ -63,8 +63,17 @@ contract RTokenConnector is Ownable, ERC20Bridge {
     function mintRToken(uint256 _mintAmount) public {
         address rTokenAddress = rToken();
         if (rTokenAddress == address(0)) return;
+
+        uint256 balanceBefore = erc20token().balanceOf(address(this));
+        uint256 rBalanceBefore = rTokenBalance();
+
         erc20token().approve(rTokenAddress, _mintAmount);
         IRToken(rTokenAddress).mint(_mintAmount);
+
+        uint256 balanceAfter = erc20token().balanceOf(address(this));
+        uint256 rBalanceAfter = rTokenBalance();
+
+        require(balanceAfter == balanceBefore.sub(_mintAmount) && rBalanceAfter == rBalanceBefore.add(_mintAmount));
     }
 
     /**
@@ -81,12 +90,26 @@ contract RTokenConnector is Ownable, ERC20Bridge {
         return addressStorage[RTOKEN];
     }
 
+    /// @dev Returns rToken balance
+    function rTokenBalance() public view returns (uint256) {
+        return IRToken(rToken()).balanceOf(address(this));
+    }
+
     /**
     * @dev Swaps rTokens to bridge's tokens
     * @param _redeemTokens Amount to swap
     */
     function _redeemRToken(uint256 _redeemTokens) internal {
         if (rToken() == address(0)) return;
+
+        uint256 balanceBefore = erc20token().balanceOf(address(this));
+        uint256 rBalanceBefore = rTokenBalance();
+
         IRToken(rToken()).redeem(_redeemTokens);
+
+        uint256 balanceAfter = erc20token().balanceOf(address(this));
+        uint256 rBalanceAfter = rTokenBalance();
+
+        require(balanceAfter == balanceBefore.add(_redeemTokens) && rBalanceAfter == rBalanceBefore.sub(_redeemTokens));
     }
 }