Change withdrawERC20 => withdrawAsset, reuse transfer logic

2 files changed, 25 insertions, 4 deletions

diff --git a/packages/contracts/test/forwarder/forwarder.ts b/packages/contracts/test/forwarder/forwarder.ts
index cd7ae59c2..28ffdeabe 100644
--- a/packages/contracts/test/forwarder/forwarder.ts
+++ b/packages/contracts/test/forwarder/forwarder.ts
@@ -36,6 +36,7 @@ describe(ContractName.Forwarder, () => {
     let feeRecipientAddress: string;
     let otherAddress: string;
     let defaultMakerAssetAddress: string;
+    let zrxAssetData: string;
 
     let weth: DummyERC20TokenContract;
     let zrxToken: DummyERC20TokenContract;
@@ -90,7 +91,7 @@ describe(ContractName.Forwarder, () => {
         erc20Wrapper.addDummyTokenContract(weth);
 
         const wethAssetData = assetDataUtils.encodeERC20AssetData(wethContract.address);
-        const zrxAssetData = assetDataUtils.encodeERC20AssetData(zrxToken.address);
+        zrxAssetData = assetDataUtils.encodeERC20AssetData(zrxToken.address);
         const exchangeInstance = await ExchangeContract.deployFrom0xArtifactAsync(
             artifacts.Exchange,
             provider,
@@ -970,6 +971,26 @@ describe(ContractName.Forwarder, () => {
             );
         });
     });
+    describe('withdrawAsset', () => {
+        it('should allow owner to withdraw ERC20 tokens', async () => {
+            const zrxWithdrawAmount = erc20Balances[forwarderContract.address][zrxToken.address];
+            await forwarderWrapper.withdrawAssetAsync(zrxAssetData, zrxWithdrawAmount, { from: owner });
+            const newBalances = await erc20Wrapper.getBalancesAsync();
+            expect(newBalances[owner][zrxToken.address]).to.be.bignumber.equal(
+                erc20Balances[owner][zrxToken.address].plus(zrxWithdrawAmount),
+            );
+            expect(newBalances[forwarderContract.address][zrxToken.address]).to.be.bignumber.equal(
+                erc20Balances[forwarderContract.address][zrxToken.address].minus(zrxWithdrawAmount),
+            );
+        });
+        it('should revert if not called by owner', async () => {
+            const zrxWithdrawAmount = erc20Balances[forwarderContract.address][zrxToken.address];
+            await expectTransactionFailedAsync(
+                forwarderWrapper.withdrawAssetAsync(zrxAssetData, zrxWithdrawAmount, { from: makerAddress }),
+                RevertReason.OnlyContractOwner,
+            );
+        });
+    });
 });
 // tslint:disable:max-file-line-count
 // tslint:enable:no-unnecessary-type-assertion
diff --git a/packages/contracts/test/utils/forwarder_wrapper.ts b/packages/contracts/test/utils/forwarder_wrapper.ts
index ef7476e36..5b9a63ddf 100644
--- a/packages/contracts/test/utils/forwarder_wrapper.ts
+++ b/packages/contracts/test/utils/forwarder_wrapper.ts
@@ -106,12 +106,12 @@ export class ForwarderWrapper {
         const tx = await this._logDecoder.getTxWithDecodedLogsAsync(txHash);
         return tx;
     }
-    public async withdrawERC20Async(
-        tokenAddress: string,
+    public async withdrawAssetAsync(
+        assetData: string,
         amount: BigNumber,
         txData: TxDataPayable,
     ): Promise<TransactionReceiptWithDecodedLogs> {
-        const txHash = await this._forwarderContract.withdrawERC20.sendTransactionAsync(tokenAddress, amount, txData);
+        const txHash = await this._forwarderContract.withdrawAsset.sendTransactionAsync(assetData, amount, txData);
         const tx = await this._logDecoder.getTxWithDecodedLogsAsync(txHash);
         return tx;
     }