Make tokenRegistry contract instantiation lazy and clear it on provider update

3 files changed, 17 insertions, 5 deletions

diff --git a/src/0x.js.ts b/src/0x.js.ts
index 506ac4187..c621a15b6 100644
--- a/src/0x.js.ts
+++ b/src/0x.js.ts
@@ -143,6 +143,7 @@ export class ZeroEx {
     public setProvider(provider: Web3.Provider) {
         this.web3Wrapper.setProvider(provider);
         this.exchange.invalidateContractInstance();
+        this.tokenRegistry.invalidateContractInstance();
     }
     /**
      * Signs an orderHash and returns it's elliptic curve signature
diff --git a/src/contract_wrappers/token_registry_wrapper.ts b/src/contract_wrappers/token_registry_wrapper.ts
index 8dd79171e..69bad2bac 100644
--- a/src/contract_wrappers/token_registry_wrapper.ts
+++ b/src/contract_wrappers/token_registry_wrapper.ts
@@ -6,17 +6,20 @@ import {ContractWrapper} from './contract_wrapper';
 import * as TokenRegistryArtifacts from '../artifacts/TokenRegistry.json';
 
 export class TokenRegistryWrapper extends ContractWrapper {
+    private tokenRegistryContractIfExists: TokenRegistryContract;
     constructor(web3Wrapper: Web3Wrapper) {
         super(web3Wrapper);
     }
+    public invalidateContractInstance() {
+        delete this.tokenRegistryContractIfExists;
+    }
     public async getTokensAsync(): Promise<Token[]> {
-        const contractInstance = await this.instantiateContractIfExistsAsync((TokenRegistryArtifacts as any));
-        const tokenRegistryContract = contractInstance as TokenRegistryContract;
+        await this.instantiateTokenRegistryContractIfDoesntExistAsync();
 
-        const addresses = await tokenRegistryContract.getTokenAddresses.call();
+        const addresses = await this.tokenRegistryContractIfExists.getTokenAddresses.call();
         const tokenMetadataPromises: Array<Promise<TokenMetadata>> = _.map(
             addresses,
-            (address: string) => (tokenRegistryContract.getTokenMetaData.call(address)),
+            (address: string) => (this.tokenRegistryContractIfExists.getTokenMetaData.call(address)),
         );
         const tokensMetadata = await Promise.all(tokenMetadataPromises);
         const tokens = _.map(tokensMetadata, metadata => {
@@ -30,4 +33,11 @@ export class TokenRegistryWrapper extends ContractWrapper {
         });
         return tokens;
     }
+    private async instantiateTokenRegistryContractIfDoesntExistAsync() {
+        if (!_.isUndefined(this.tokenRegistryContractIfExists)) {
+            return;
+        }
+        const contractInstance = await this.instantiateContractIfExistsAsync((TokenRegistryArtifacts as any));
+        this.tokenRegistryContractIfExists = contractInstance as TokenRegistryContract;
+    }
 }
diff --git a/test/0x.js_test.ts b/test/0x.js_test.ts
index 1e76ae71a..4125b83fe 100644
--- a/test/0x.js_test.ts
+++ b/test/0x.js_test.ts
@@ -26,8 +26,9 @@ describe('ZeroEx library', () => {
             (newProvider as any).zeroExTestId = 1;
             zeroEx.setProvider(newProvider);
 
-            // Check that exchangeContract instance removed after provider update
+            // Check that contractInstances with old provider are removed after provider update
             expect((zeroEx.exchange as any).exchangeContractIfExists).to.be.an('undefined');
+            expect((zeroEx.tokenRegistry as any).tokenRegistryContractIfExists).to.be.an('undefined');
 
             // Check that all nested web3 instances return the updated provider
             const nestedWeb3WrapperProvider = (zeroEx as any).web3Wrapper.getCurrentProvider();