- Refactor assetWrapper to contain more of the normalizing logic instead of erc20Wrapper and erc721Wrapper

- Add burn method to DummyERC721Token
- Add additional methods to assetWrapper to set balance/allowances on ERC20 and ERC721 tokens
- Use approve instead of approveAll for ERC721 tokens

4 files changed, 295 insertions, 67 deletions

diff --git a/packages/contracts/src/utils/asset_wrapper.ts b/packages/contracts/src/utils/asset_wrapper.ts
index 4c345aa30..462a5086a 100644
--- a/packages/contracts/src/utils/asset_wrapper.ts
+++ b/packages/contracts/src/utils/asset_wrapper.ts
@@ -1,9 +1,13 @@
 import { assetProxyUtils } from '@0xproject/order-utils';
-import { BigNumber } from '@0xproject/utils';
+import { BigNumber, errorUtils } from '@0xproject/utils';
 import * as _ from 'lodash';
 
 import { AbstractAssetWrapper } from '../abstract/abstract_asset_wrapper';
 
+import { constants } from './constants';
+import { ERC20Wrapper } from './erc20_wrapper';
+import { ERC721Wrapper } from './erc721_wrapper';
+
 interface ProxyIdToAssetWrappers {
     [proxyId: number]: AbstractAssetWrapper;
 }
@@ -17,16 +21,179 @@ export class AssetWrapper {
             this._proxyIdToAssetWrappers[proxyId] = assetWrapper;
         });
     }
-    public async getBalanceAsync(owner: string, assetData: string): Promise<BigNumber> {
+    public async getBalanceAsync(userAddress: string, assetData: string): Promise<BigNumber> {
+        const proxyId = assetProxyUtils.decodeAssetDataId(assetData);
+        switch (proxyId) {
+            case constants.ERC20_PROXY_ID: {
+                const assetWrapper = this._proxyIdToAssetWrappers[proxyId] as ERC20Wrapper;
+                const balance = await assetWrapper.getBalanceAsync(userAddress, assetData);
+                return balance;
+            }
+
+            case constants.ERC721_PROXY_ID: {
+                const assetWrapper = this._proxyIdToAssetWrappers[proxyId] as ERC721Wrapper;
+                const assetProxyData = assetProxyUtils.decodeERC721AssetData(assetData);
+                const isOwner = await assetWrapper.isOwnerAsync(
+                    userAddress,
+                    assetProxyData.tokenAddress,
+                    assetProxyData.tokenId,
+                );
+                const balance = isOwner ? new BigNumber(1) : new BigNumber(0);
+                return balance;
+            }
+
+            default:
+                throw errorUtils.spawnSwitchErr('proxyId', proxyId);
+        }
+    }
+    public async setBalanceAsync(userAddress: string, assetData: string, desiredBalance: BigNumber): Promise<void> {
+        const proxyId = assetProxyUtils.decodeAssetDataId(assetData);
+        switch (proxyId) {
+            case constants.ERC20_PROXY_ID: {
+                const assetWrapper = this._proxyIdToAssetWrappers[proxyId] as ERC20Wrapper;
+                await assetWrapper.setBalanceAsync(userAddress, assetData, desiredBalance);
+                return;
+            }
+
+            case constants.ERC721_PROXY_ID: {
+                if (!desiredBalance.eq(0) && !desiredBalance.eq(1)) {
+                    throw new Error(`Balance for ERC721 token can only be set to 0 or 1. Got: ${desiredBalance}`);
+                }
+                const erc721Wrapper = this._proxyIdToAssetWrappers[proxyId] as ERC721Wrapper;
+                const assetProxyData = assetProxyUtils.decodeERC721AssetData(assetData);
+                const doesTokenExist = erc721Wrapper.doesTokenExistAsync(
+                    assetProxyData.tokenAddress,
+                    assetProxyData.tokenId,
+                );
+                if (!doesTokenExist && desiredBalance.eq(1)) {
+                    await erc721Wrapper.mintAsync(assetProxyData.tokenAddress, assetProxyData.tokenId, userAddress);
+                    return;
+                } else if (!doesTokenExist && desiredBalance.eq(0)) {
+                    return; // noop
+                }
+                const tokenOwner = await erc721Wrapper.ownerOfAsync(
+                    assetProxyData.tokenAddress,
+                    assetProxyData.tokenId,
+                );
+                if (userAddress !== tokenOwner && desiredBalance.eq(1)) {
+                    await erc721Wrapper.transferFromAsync(
+                        assetProxyData.tokenAddress,
+                        assetProxyData.tokenId,
+                        tokenOwner,
+                        userAddress,
+                    );
+                } else if (
+                    (userAddress !== tokenOwner && desiredBalance.eq(0)) ||
+                    (tokenOwner === userAddress && desiredBalance.eq(1))
+                ) {
+                    return; // noop
+                } else if (tokenOwner === userAddress && desiredBalance.eq(0)) {
+                    // Burn token
+                    await erc721Wrapper.burnAsync(assetProxyData.tokenAddress, assetProxyData.tokenId, userAddress);
+                    return;
+                }
+                break;
+            }
+
+            default:
+                throw errorUtils.spawnSwitchErr('proxyId', proxyId);
+        }
+    }
+    public async getProxyAllowanceAsync(userAddress: string, assetData: string): Promise<BigNumber> {
         const proxyId = assetProxyUtils.decodeAssetDataId(assetData);
-        const assetWrapper = this._proxyIdToAssetWrappers[proxyId];
-        const balance = await assetWrapper.getBalanceAsync(owner, assetData);
-        return balance;
+        switch (proxyId) {
+            case constants.ERC20_PROXY_ID: {
+                const assetWrapper = this._proxyIdToAssetWrappers[proxyId] as ERC20Wrapper;
+                const allowance = await assetWrapper.getProxyAllowanceAsync(userAddress, assetData);
+                return allowance;
+            }
+
+            case constants.ERC721_PROXY_ID: {
+                const assetWrapper = this._proxyIdToAssetWrappers[proxyId] as ERC721Wrapper;
+                const erc721ProxyData = assetProxyUtils.decodeERC721AssetData(assetData);
+                const isProxyApproved = await assetWrapper.isProxyApprovedAsync(
+                    erc721ProxyData.tokenAddress,
+                    erc721ProxyData.tokenId,
+                );
+                const isProxyApprovedForAllAsync = await assetWrapper.isProxyApprovedForAllAsync(
+                    userAddress,
+                    erc721ProxyData.tokenAddress,
+                    erc721ProxyData.tokenId,
+                );
+                const allowance = isProxyApproved || isProxyApprovedForAllAsync ? new BigNumber(1) : new BigNumber(0);
+                return allowance;
+            }
+
+            default:
+                throw errorUtils.spawnSwitchErr('proxyId', proxyId);
+        }
     }
-    public async getProxyAllowanceAsync(owner: string, assetData: string): Promise<BigNumber> {
+    public async setProxyAllowanceAsync(
+        userAddress: string,
+        assetData: string,
+        desiredAllowance: BigNumber,
+    ): Promise<void> {
         const proxyId = assetProxyUtils.decodeAssetDataId(assetData);
-        const assetWrapper = this._proxyIdToAssetWrappers[proxyId];
-        const balance = await assetWrapper.getProxyAllowanceAsync(owner, assetData);
-        return balance;
+        switch (proxyId) {
+            case constants.ERC20_PROXY_ID: {
+                const assetWrapper = this._proxyIdToAssetWrappers[proxyId] as ERC20Wrapper;
+                await assetWrapper.setAllowanceAsync(userAddress, assetData, desiredAllowance);
+                return;
+            }
+
+            case constants.ERC721_PROXY_ID: {
+                if (!desiredAllowance.eq(0) && !desiredAllowance.eq(1)) {
+                    throw new Error(`Allowance for ERC721 token can only be set to 0 or 1. Got: ${desiredAllowance}`);
+                }
+                const erc721Wrapper = this._proxyIdToAssetWrappers[proxyId] as ERC721Wrapper;
+                const assetProxyData = assetProxyUtils.decodeERC721AssetData(assetData);
+
+                const doesTokenExist = await erc721Wrapper.doesTokenExistAsync(
+                    assetProxyData.tokenAddress,
+                    assetProxyData.tokenId,
+                );
+                if (!doesTokenExist) {
+                    throw new Error(
+                        `Cannot setProxyAllowance on non-existent token: ${assetProxyData.tokenAddress} ${
+                            assetProxyData.tokenId
+                        }`,
+                    );
+                }
+                const isProxyApprovedForAll = await erc721Wrapper.isProxyApprovedForAllAsync(
+                    userAddress,
+                    assetProxyData.tokenAddress,
+                    assetProxyData.tokenId,
+                );
+                // TODO: We should have a way to deal with this. Things get hairier once we are testing
+                // batch fills
+                if (isProxyApprovedForAll) {
+                    throw new Error(`We don't currently support the use of "approveAll" functionality for ERC721.`);
+                }
+
+                const isProxyApproved = await erc721Wrapper.isProxyApprovedAsync(
+                    assetProxyData.tokenAddress,
+                    assetProxyData.tokenId,
+                );
+                if (!isProxyApproved && desiredAllowance.eq(1)) {
+                    await erc721Wrapper.approveProxyAsync(assetProxyData.tokenAddress, assetProxyData.tokenId);
+                } else if (isProxyApproved && desiredAllowance.eq(0)) {
+                    await erc721Wrapper.approveAsync(
+                        constants.NULL_ADDRESS,
+                        assetProxyData.tokenAddress,
+                        assetProxyData.tokenId,
+                    );
+                } else if (
+                    (!isProxyApproved && desiredAllowance.eq(0)) ||
+                    (isProxyApproved && desiredAllowance.eq(1))
+                ) {
+                    return; // noop
+                }
+
+                break;
+            }
+
+            default:
+                throw errorUtils.spawnSwitchErr('proxyId', proxyId);
+        }
     }
 }
diff --git a/packages/contracts/src/utils/constants.ts b/packages/contracts/src/utils/constants.ts
index ec3c8fd36..c8b684c41 100644
--- a/packages/contracts/src/utils/constants.ts
+++ b/packages/contracts/src/utils/constants.ts
@@ -27,6 +27,8 @@ export const constants = {
     LIB_BYTES_GREATER_OR_EQUAL_TO_SOURCE_BYTES_LENGTH_REQUIRED: 'GREATER_OR_EQUAL_TO_SOURCE_BYTES_LENGTH_REQUIRED',
     ERC20_INSUFFICIENT_BALANCE: 'Insufficient balance to complete transfer.',
     ERC20_INSUFFICIENT_ALLOWANCE: 'Insufficient allowance to complete transfer.',
+    ERC20_PROXY_ID: 1,
+    ERC721_PROXY_ID: 2,
     TESTRPC_NETWORK_ID: 50,
     // Note(albrow): In practice V8 and most other engines limit the minimum
     // interval for setInterval to 10ms. We still set it to 0 here in order to
diff --git a/packages/contracts/src/utils/erc20_wrapper.ts b/packages/contracts/src/utils/erc20_wrapper.ts
index 2d367ae1c..ca2ca492b 100644
--- a/packages/contracts/src/utils/erc20_wrapper.ts
+++ b/packages/contracts/src/utils/erc20_wrapper.ts
@@ -83,29 +83,30 @@ export class ERC20Wrapper {
             }
         }
     }
-    public async getBalanceAsync(owner: string, assetData: string): Promise<BigNumber> {
-        const erc20ProxyData = assetProxyUtils.decodeERC20AssetData(assetData);
-        const tokenAddress = erc20ProxyData.tokenAddress;
-        const tokenContractIfExists = _.find(this._dummyTokenContracts, c => c.address === tokenAddress);
-        if (_.isUndefined(tokenContractIfExists)) {
-            throw new Error(`Token: ${tokenAddress} was not deployed through ERC20Wrapper`);
-        }
-        const balance = new BigNumber(await tokenContractIfExists.balanceOf.callAsync(owner));
+    public async getBalanceAsync(userAddress: string, assetData: string): Promise<BigNumber> {
+        const tokenContract = this._getTokenContractFromAssetData(assetData);
+        const balance = new BigNumber(await tokenContract.balanceOf.callAsync(userAddress));
         return balance;
     }
-    public async getProxyAllowanceAsync(owner: string, assetData: string): Promise<BigNumber> {
-        const erc20ProxyData = assetProxyUtils.decodeERC20AssetData(assetData);
-        const tokenAddress = erc20ProxyData.tokenAddress;
-        this._validateProxyContractExistsOrThrow();
-        const tokenContractIfExists = _.find(this._dummyTokenContracts, c => c.address === tokenAddress);
-        if (_.isUndefined(tokenContractIfExists)) {
-            throw new Error(`Token: ${tokenAddress} was not deployed through ERC20Wrapper`);
-        }
-        const allowance = new BigNumber(
-            await tokenContractIfExists.allowance.callAsync(owner, (this._proxyContract as ERC20ProxyContract).address),
-        );
+    public async setBalanceAsync(userAddress: string, assetData: string, amount: BigNumber): Promise<void> {
+        const tokenContract = this._getTokenContractFromAssetData(assetData);
+        await tokenContract.setBalance.sendTransactionAsync(userAddress, amount, {
+            from: this._contractOwnerAddress,
+        });
+    }
+    public async getProxyAllowanceAsync(userAddress: string, assetData: string): Promise<BigNumber> {
+        const tokenContract = this._getTokenContractFromAssetData(assetData);
+        const proxyAddress = (this._proxyContract as ERC20ProxyContract).address;
+        const allowance = new BigNumber(await tokenContract.allowance.callAsync(userAddress, proxyAddress));
         return allowance;
     }
+    public async setAllowanceAsync(userAddress: string, assetData: string, amount: BigNumber): Promise<void> {
+        const tokenContract = this._getTokenContractFromAssetData(assetData);
+        const proxyAddress = (this._proxyContract as ERC20ProxyContract).address;
+        await tokenContract.approve.sendTransactionAsync(proxyAddress, amount, {
+            from: userAddress,
+        });
+    }
     public async getBalancesAsync(): Promise<ERC20BalancesByOwner> {
         this._validateDummyTokenContractsExistOrThrow();
         const balancesByOwner: ERC20BalancesByOwner = {};
@@ -138,6 +139,15 @@ export class ERC20Wrapper {
         const tokenAddresses = _.map(this._dummyTokenContracts, dummyTokenContract => dummyTokenContract.address);
         return tokenAddresses;
     }
+    private _getTokenContractFromAssetData(assetData: string): DummyERC20TokenContract {
+        const erc20ProxyData = assetProxyUtils.decodeERC20AssetData(assetData);
+        const tokenAddress = erc20ProxyData.tokenAddress;
+        const tokenContractIfExists = _.find(this._dummyTokenContracts, c => c.address === tokenAddress);
+        if (_.isUndefined(tokenContractIfExists)) {
+            throw new Error(`Token: ${tokenAddress} was not deployed through ERC20Wrapper`);
+        }
+        return tokenContractIfExists;
+    }
     private _validateDummyTokenContractsExistOrThrow(): void {
         if (_.isUndefined(this._dummyTokenContracts)) {
             throw new Error('Dummy ERC20 tokens not yet deployed, please call "deployDummyTokensAsync"');
diff --git a/packages/contracts/src/utils/erc721_wrapper.ts b/packages/contracts/src/utils/erc721_wrapper.ts
index 6a6f10fa7..7da3b6781 100644
--- a/packages/contracts/src/utils/erc721_wrapper.ts
+++ b/packages/contracts/src/utils/erc721_wrapper.ts
@@ -63,12 +63,7 @@ export class ERC721Wrapper {
             for (const tokenOwnerAddress of this._tokenOwnerAddresses) {
                 for (let i = 0; i < constants.NUM_ERC721_TOKENS_TO_MINT; i++) {
                     const tokenId = generatePseudoRandomSalt();
-                    await this._web3Wrapper.awaitTransactionSuccessAsync(
-                        await dummyTokenContract.mint.sendTransactionAsync(tokenOwnerAddress, tokenId, {
-                            from: this._contractOwnerAddress,
-                        }),
-                        constants.AWAIT_TRANSACTION_MINED_MS,
-                    );
+                    await this.mintAsync(dummyTokenContract.address, tokenId, tokenOwnerAddress);
                     if (_.isUndefined(this._initialTokenIdsByOwner[tokenOwnerAddress])) {
                         this._initialTokenIdsByOwner[tokenOwnerAddress] = {
                             [dummyTokenContract.address]: [],
@@ -78,45 +73,92 @@ export class ERC721Wrapper {
                         this._initialTokenIdsByOwner[tokenOwnerAddress][dummyTokenContract.address] = [];
                     }
                     this._initialTokenIdsByOwner[tokenOwnerAddress][dummyTokenContract.address].push(tokenId);
+
+                    await this.approveProxyAsync(dummyTokenContract.address, tokenId);
                 }
-                const shouldApprove = true;
-                await this._web3Wrapper.awaitTransactionSuccessAsync(
-                    await dummyTokenContract.setApprovalForAll.sendTransactionAsync(
-                        (this._proxyContract as ERC721ProxyContract).address,
-                        shouldApprove,
-                        { from: tokenOwnerAddress },
-                    ),
-                    constants.AWAIT_TRANSACTION_MINED_MS,
-                );
             }
         }
     }
-    public async getBalanceAsync(owner: string, assetData: string): Promise<BigNumber> {
-        const erc721ProxyData = assetProxyUtils.decodeERC721AssetData(assetData);
-        const tokenAddress = erc721ProxyData.tokenAddress;
-        const tokenContractIfExists = _.find(this._dummyTokenContracts, c => c.address === tokenAddress);
-        if (_.isUndefined(tokenContractIfExists)) {
-            throw new Error(`Token: ${tokenAddress} was not deployed through ERC20Wrapper`);
-        }
-        const tokenId = erc721ProxyData.tokenId;
-        const tokenOwner = await tokenContractIfExists.ownerOf.callAsync(tokenId);
-        const balance = tokenOwner === owner ? new BigNumber(1) : new BigNumber(0);
-        return balance;
-    }
-    public async getProxyAllowanceAsync(owner: string, assetData: string): Promise<BigNumber> {
-        const erc721ProxyData = assetProxyUtils.decodeERC721AssetData(assetData);
-        const tokenAddress = erc721ProxyData.tokenAddress;
+    public async doesTokenExistAsync(tokenAddress: string, tokenId: BigNumber): Promise<boolean> {
+        const tokenContract = this._getTokenContractFromAssetData(tokenAddress);
+        const doesExist = await tokenContract.exists.callAsync(tokenId);
+        return doesExist;
+    }
+    public async approveProxyAsync(tokenAddress: string, tokenId: BigNumber): Promise<void> {
+        const proxyAddress = (this._proxyContract as ERC721ProxyContract).address;
+        await this.approveAsync(proxyAddress, tokenAddress, tokenId);
+    }
+    public async approveAsync(to: string, tokenAddress: string, tokenId: BigNumber): Promise<void> {
+        const tokenContract = this._getTokenContractFromAssetData(tokenAddress);
+        const tokenOwner = await this.ownerOfAsync(tokenAddress, tokenId);
+        await this._web3Wrapper.awaitTransactionSuccessAsync(
+            await tokenContract.approve.sendTransactionAsync(to, tokenId, {
+                from: tokenOwner,
+            }),
+            constants.AWAIT_TRANSACTION_MINED_MS,
+        );
+    }
+    public async transferFromAsync(
+        tokenAddress: string,
+        tokenId: BigNumber,
+        currentOwner: string,
+        userAddress: string,
+    ): Promise<void> {
+        const tokenContract = this._getTokenContractFromAssetData(tokenAddress);
+        await this._web3Wrapper.awaitTransactionSuccessAsync(
+            await tokenContract.transferFrom.sendTransactionAsync(currentOwner, userAddress, tokenId, {
+                from: currentOwner,
+            }),
+            constants.AWAIT_TRANSACTION_MINED_MS,
+        );
+    }
+    public async mintAsync(tokenAddress: string, tokenId: BigNumber, userAddress: string): Promise<void> {
+        const tokenContract = this._getTokenContractFromAssetData(tokenAddress);
+        await this._web3Wrapper.awaitTransactionSuccessAsync(
+            await tokenContract.mint.sendTransactionAsync(userAddress, tokenId, {
+                from: this._contractOwnerAddress,
+            }),
+            constants.AWAIT_TRANSACTION_MINED_MS,
+        );
+    }
+    public async burnAsync(tokenAddress: string, tokenId: BigNumber, owner: string): Promise<void> {
+        const tokenContract = this._getTokenContractFromAssetData(tokenAddress);
+        await this._web3Wrapper.awaitTransactionSuccessAsync(
+            await tokenContract.burn.sendTransactionAsync(owner, tokenId, {
+                from: this._contractOwnerAddress,
+            }),
+            constants.AWAIT_TRANSACTION_MINED_MS,
+        );
+    }
+    public async ownerOfAsync(tokenAddress: string, tokenId: BigNumber): Promise<string> {
+        const tokenContract = this._getTokenContractFromAssetData(tokenAddress);
+        const owner = await tokenContract.ownerOf.callAsync(tokenId);
+        return owner;
+    }
+    public async isOwnerAsync(userAddress: string, tokenAddress: string, tokenId: BigNumber): Promise<boolean> {
+        const tokenContract = this._getTokenContractFromAssetData(tokenAddress);
+        const tokenOwner = await tokenContract.ownerOf.callAsync(tokenId);
+        const isOwner = tokenOwner === userAddress;
+        return isOwner;
+    }
+    public async isProxyApprovedForAllAsync(
+        userAddress: string,
+        tokenAddress: string,
+        tokenId: BigNumber,
+    ): Promise<boolean> {
         this._validateProxyContractExistsOrThrow();
-        const tokenContractIfExists = _.find(this._dummyTokenContracts, c => c.address === tokenAddress);
-        if (_.isUndefined(tokenContractIfExists)) {
-            throw new Error(`Token: ${tokenAddress} was not deployed through ERC20Wrapper`);
-        }
+        const tokenContract = this._getTokenContractFromAssetData(tokenAddress);
         const operator = (this._proxyContract as ERC721ProxyContract).address;
-        const didApproveAll = await tokenContractIfExists.isApprovedForAll.callAsync(owner, operator);
-        const tokenId = erc721ProxyData.tokenId;
-        const allowedAddress = await tokenContractIfExists.getApproved.callAsync(tokenId);
-        const allowance = allowedAddress === operator || didApproveAll ? new BigNumber(1) : new BigNumber(0);
-        return allowance;
+        const didApproveAll = await tokenContract.isApprovedForAll.callAsync(userAddress, operator);
+        return didApproveAll;
+    }
+    public async isProxyApprovedAsync(tokenAddress: string, tokenId: BigNumber): Promise<boolean> {
+        this._validateProxyContractExistsOrThrow();
+        const tokenContract = this._getTokenContractFromAssetData(tokenAddress);
+        const approvedAddress = await tokenContract.getApproved.callAsync(tokenId);
+        const proxyAddress = (this._proxyContract as ERC721ProxyContract).address;
+        const isProxyAnApprovedOperator = approvedAddress === proxyAddress;
+        return isProxyAnApprovedOperator;
     }
     public async getBalancesAsync(): Promise<ERC721TokenIdsByOwner> {
         this._validateDummyTokenContractsExistOrThrow();
@@ -160,6 +202,13 @@ export class ERC721Wrapper {
         const tokenAddresses = _.map(this._dummyTokenContracts, dummyTokenContract => dummyTokenContract.address);
         return tokenAddresses;
     }
+    private _getTokenContractFromAssetData(tokenAddress: string): DummyERC721TokenContract {
+        const tokenContractIfExists = _.find(this._dummyTokenContracts, c => c.address === tokenAddress);
+        if (_.isUndefined(tokenContractIfExists)) {
+            throw new Error(`Token: ${tokenAddress} was not deployed through ERC20Wrapper`);
+        }
+        return tokenContractIfExists;
+    }
     private _validateDummyTokenContractsExistOrThrow(): void {
         if (_.isUndefined(this._dummyTokenContracts)) {
             throw new Error('Dummy ERC721 tokens not yet deployed, please call "deployDummyTokensAsync"');