From a24ee929fda41a0067ec56e8d7f0c7a39949f89e Mon Sep 17 00:00:00 2001
From: Lecky Lao <leckylao@gmail.com>
Date: Wed, 1 Jul 2020 00:17:20 +1000
Subject: [PATCH] added method withdraw and test;

---
 contracts/connectors/curvesbtc.sol | 48 ++++++++++++++++++++++++++++++
 test/CurveSBTCProtocol.js          | 25 +++++++++++++---
 2 files changed, 69 insertions(+), 4 deletions(-)

diff --git a/contracts/connectors/curvesbtc.sol b/contracts/connectors/curvesbtc.sol
index e796188..f599d86 100644
--- a/contracts/connectors/curvesbtc.sol
+++ b/contracts/connectors/curvesbtc.sol
@@ -156,4 +156,52 @@ contract CurveSBTCProtocol is CurveSBTCHelpers {
     emitEvent(_eventCode, _eventParam);
   }
 
+  /**
+  * @dev Withdraw Token.
+  * @param token token address.
+  * @param amt token amount.
+  * @param unitAmt unit amount of curve_amt/token_amt with slippage.
+  * @param getId Get token amount at this ID from `InstaMemory` Contract.
+  * @param setId Set token amount at this ID in `InstaMemory` Contract.
+  */
+  function withdraw(
+    address token,
+    uint256 amt,
+    uint256 unitAmt,
+    uint getId,
+    uint setId
+  ) external payable {
+    uint _amt = getUint(getId, amt);
+    int128 tokenId = getTokenI(token);
+
+    ERC20 curveTokenContract = ERC20(getCurveTokenAddr());
+    ICurve curveSwap = ICurve(getCurveSwapAddr());
+
+    uint _curveAmt;
+    uint[3] memory _amts;
+    if (_amt == uint(-1)) {
+      _curveAmt = curveTokenContract.balanceOf(address(this));
+      _amt = curveSwap.calc_withdraw_one_coin(_curveAmt, tokenId);
+      _amts[uint(tokenId)] = _amt;
+    } else {
+      _amts[uint(tokenId)] = _amt;
+      _curveAmt = curveSwap.calc_token_amount(_amts, false);
+    }
+
+    uint _amt18 = convertTo18(ERC20(token).decimals(), _amt);
+    uint _slippageAmt = wmul(unitAmt, _amt18);
+
+    curveTokenContract.approve(address(curveSwap), 0);
+    curveTokenContract.approve(address(curveSwap), _slippageAmt);
+
+    curveSwap.remove_liquidity_imbalance(_amts, _slippageAmt);
+
+    setUint(setId, _amt);
+
+    emit LogWithdraw(token, _amt, _curveAmt, getId, setId);
+    bytes32 _eventCode = keccak256("LogWithdraw(address,uint256,uint256,uint256,uint256)");
+    bytes memory _eventParam = abi.encode(token, _amt, _curveAmt, getId, setId);
+    emitEvent(_eventCode, _eventParam);
+  }
+
 }
diff --git a/test/CurveSBTCProtocol.js b/test/CurveSBTCProtocol.js
index adb3071..83c820e 100644
--- a/test/CurveSBTCProtocol.js
+++ b/test/CurveSBTCProtocol.js
@@ -82,7 +82,7 @@ contract('CurveSBTCProtocol', async accounts => {
       "0x075b1bb99792c9e1041ba13afef80c91a1e70fb3"
     )
 
-    const tx = await contract.deposit(
+    const txDeposit = await contract.deposit(
       erc20.wbtc.address,
       10000000,
       ( 0.09 / 0.1 * 1e18 ).toString(),
@@ -93,11 +93,28 @@ contract('CurveSBTCProtocol', async accounts => {
         from: sender
       }
     );
-    console.log(tx);
+    console.log(txDeposit);
 
-    const balance = await curveTokenContract.methods.balanceOf(sender);
+    const balanceDeposit = await curveTokenContract.methods.balanceOf(sender);
 
-    expect(balance).to.be.at.least(ether("0.09"));
+    expect(balanceDeposit).to.be.at.least(ether("0.09"));
 
+    const txWithdraw = await contract.withdraw(
+      erc20.wbtc.address,
+      10000000,
+      ( 0.09 / 0.1 * 1e18 ).toString(),
+      0,
+      0,
+      {
+        gas: 4000000,
+        from: sender
+      }
+    );
+    console.log(txWithdraw);
+
+    const balanceWithdraw = await curveTokenContract.methods.balanceOf(sender);
+
+    expect(balanceWithdraw).to.equal(0);
   });
+
 });