1 files changed, 57 insertions, 27 deletions

diff --git a/packages/migrations/src/migration.ts b/packages/migrations/src/migration.ts
index 99d1719f1..3d97409e6 100644
--- a/packages/migrations/src/migration.ts
+++ b/packages/migrations/src/migration.ts
@@ -68,13 +68,46 @@ export async function runMigrationsAsync(provider: Provider, txDefaults: Partial
         secondsRequired,
     );
 
+    // Dummy ERC20 tokens
+    for (const token of erc20TokenInfo) {
+        const totalSupply = new BigNumber(1000000000000000000000000000);
+        // tslint:disable-next-line:no-unused-variable
+        const dummyErc20Token = await wrappers.DummyERC20TokenContract.deployFrom0xArtifactAsync(
+            artifacts.DummyERC20Token,
+            provider,
+            txDefaults,
+            token.name,
+            token.symbol,
+            token.decimals,
+            totalSupply,
+        );
+    }
+
+    // ERC721
+    // tslint:disable-next-line:no-unused-variable
+    const cryptoKittieToken = await wrappers.DummyERC721TokenContract.deployFrom0xArtifactAsync(
+        artifacts.DummyERC721Token,
+        provider,
+        txDefaults,
+        erc721TokenInfo[0].name,
+        erc721TokenInfo[0].symbol,
+    );
+
+    const multiAssetProxy = await wrappers.MultiAssetProxyContract.deployFrom0xArtifactAsync(
+        artifacts.MultiAssetProxy,
+        provider,
+        txDefaults,
+    );
+
+    // Note: Perform all setup after contract deployments to have consistent addresses
+    // available in the snapshots
     await web3Wrapper.awaitTransactionSuccessAsync(
         await erc20Proxy.addAuthorizedAddress.sendTransactionAsync(exchange.address, {
             from: owner,
         }),
     );
     await web3Wrapper.awaitTransactionSuccessAsync(
-        await erc20Proxy.transferOwnership.sendTransactionAsync(assetProxyOwner.address, {
+        await erc20Proxy.addAuthorizedAddress.sendTransactionAsync(multiAssetProxy.address, {
             from: owner,
         }),
     );
@@ -84,7 +117,7 @@ export async function runMigrationsAsync(provider: Provider, txDefaults: Partial
         }),
     );
     await web3Wrapper.awaitTransactionSuccessAsync(
-        await erc721Proxy.transferOwnership.sendTransactionAsync(assetProxyOwner.address, {
+        await erc721Proxy.addAuthorizedAddress.sendTransactionAsync(multiAssetProxy.address, {
             from: owner,
         }),
     );
@@ -96,30 +129,8 @@ export async function runMigrationsAsync(provider: Provider, txDefaults: Partial
     await web3Wrapper.awaitTransactionSuccessAsync(
         await exchange.registerAssetProxy.sendTransactionAsync(erc721Proxy.address),
     );
-
-    // Dummy ERC20 tokens
-    for (const token of erc20TokenInfo) {
-        const totalSupply = new BigNumber(1000000000000000000000000000);
-        // tslint:disable-next-line:no-unused-variable
-        const dummyErc20Token = await wrappers.DummyERC20TokenContract.deployFrom0xArtifactAsync(
-            artifacts.DummyERC20Token,
-            provider,
-            txDefaults,
-            token.name,
-            token.symbol,
-            token.decimals,
-            totalSupply,
-        );
-    }
-
-    // ERC721
-    // tslint:disable-next-line:no-unused-variable
-    const cryptoKittieToken = await wrappers.DummyERC721TokenContract.deployFrom0xArtifactAsync(
-        artifacts.DummyERC721Token,
-        provider,
-        txDefaults,
-        erc721TokenInfo[0].name,
-        erc721TokenInfo[0].symbol,
+    await web3Wrapper.awaitTransactionSuccessAsync(
+        await exchange.registerAssetProxy.sendTransactionAsync(multiAssetProxy.address),
     );
 
     // Forwarder
@@ -149,6 +160,23 @@ export async function runMigrationsAsync(provider: Provider, txDefaults: Partial
         exchange.address,
     );
 
+    // Transfer Ownership to the Asset Proxy Owner
+    await web3Wrapper.awaitTransactionSuccessAsync(
+        await erc20Proxy.transferOwnership.sendTransactionAsync(assetProxyOwner.address, {
+            from: owner,
+        }),
+    );
+    await web3Wrapper.awaitTransactionSuccessAsync(
+        await erc721Proxy.transferOwnership.sendTransactionAsync(assetProxyOwner.address, {
+            from: owner,
+        }),
+    );
+    await web3Wrapper.awaitTransactionSuccessAsync(
+        await multiAssetProxy.transferOwnership.sendTransactionAsync(assetProxyOwner.address, {
+            from: owner,
+        }),
+    );
+
     // Fund the Forwarder with ZRX
     const zrxDecimals = await zrxToken.decimals.callAsync();
     const zrxForwarderAmount = Web3Wrapper.toBaseUnitAmount(new BigNumber(5000), zrxDecimals);
@@ -156,7 +184,7 @@ export async function runMigrationsAsync(provider: Provider, txDefaults: Partial
         await zrxToken.transfer.sendTransactionAsync(forwarder.address, zrxForwarderAmount, txDefaults),
     );
 
-    return {
+    const contractAddresses = {
         erc20Proxy: erc20Proxy.address,
         erc721Proxy: erc721Proxy.address,
         zrxToken: zrxToken.address,
@@ -167,6 +195,8 @@ export async function runMigrationsAsync(provider: Provider, txDefaults: Partial
         orderValidator: orderValidator.address,
         dutchAuction: dutchAuction.address,
     };
+
+    return contractAddresses;
 }
 
 let _cachedContractAddresses: ContractAddresses;