detect tokens when submit password and new account selected

1 files changed, 39 insertions, 13 deletions

diff --git a/app/scripts/controllers/detect-tokens.js b/app/scripts/controllers/detect-tokens.js
index e245a7f9b..db21f7489 100644
--- a/app/scripts/controllers/detect-tokens.js
+++ b/app/scripts/controllers/detect-tokens.js
@@ -1,9 +1,7 @@
 const Web3 = require('web3')
 const contracts = require('eth-contract-metadata')
 const { warn } = require('loglevel')
-const {
-    MAINNET,
-    } = require('./network/enums')
+const { MAINNET } = require('./network/enums')
 // By default, poll every 3 minutes
 const DEFAULT_INTERVAL = 180 * 1000
 const ERC20_ABI = [{'constant': true, 'inputs': [{'name': '_owner', 'type': 'address'}], 'name': 'balanceOf', 'outputs': [{'name': 'balance', 'type': 'uint256'}], 'payable': false, 'type': 'function'}]
@@ -22,14 +20,15 @@ class DetectTokensController {
     this.preferences = preferences
     this.interval = interval
     this.network = network
+    this._isActive = false
   }
 
-   /**
+  /**
    * For each token in eth-contract-metada, find check selectedAddress balance.
    *
    */
-  async exploreNewTokens () {
-    if (!this.isActive) { return }
+  async detectNewTokens () {
+    if (!this._isActive) { return }
     if (this._network.store.getState().provider.type !== MAINNET) { return }
     this.web3.setProvider(this._network._provider)
     for (const contractAddress in contracts) {
@@ -60,27 +59,43 @@ class DetectTokensController {
   }
 
   /**
+   * Restart token detection polling period and call detectNewTokens
+   * in case of address change or user session initialization.
+   *
+   */
+  restartTokenDetection () {
+    if (this._isActive && this.selectedAddress) {
+      this.detectNewTokens()
+      this.interval = DEFAULT_INTERVAL
+    }
+  }
+
+  /**
    * @type {Number}
    */
   set interval (interval) {
     this._handle && clearInterval(this._handle)
     if (!interval) { return }
-    this._handle = setInterval(() => { this.exploreNewTokens() }, interval)
+    this._handle = setInterval(() => { this.detectNewTokens() }, interval)
   }
 
-    /**
+  /**
+   * In setter when selectedAddress is changed, detectNewTokens and restart polling
    * @type {Object}
    */
   set preferences (preferences) {
     if (!preferences) { return }
     this._preferences = preferences
-    this.tokenAddresses = preferences.store.getState().tokens.map((obj) => { return obj.address })
-    this.selectedAddress = preferences.store.getState().selectedAddress
-    preferences.store.subscribe(({ tokens = [] }) => { this.tokenAddresses = tokens.map((obj) => { return obj.address }) })
-    preferences.store.subscribe(({ selectedAddress = [] }) => { this.selectedAddress = selectedAddress })
+    preferences.store.subscribe(({ tokens }) => { this.tokenAddresses = tokens.map((obj) => { return obj.address }) })
+    preferences.store.subscribe(({ selectedAddress }) => {
+      if (this.selectedAddress !== selectedAddress) {
+        this.selectedAddress = selectedAddress
+        this.restartTokenDetection()
+      }
+    })
   }
 
-    /**
+  /**
    * @type {Object}
    */
   set network (network) {
@@ -88,6 +103,17 @@ class DetectTokensController {
     this._network = network
     this.web3 = new Web3(network._provider)
   }
+
+  /**
+  * In setter, when _isActive is changed, detectNewTokens and restart polling
+  * @type {Object}
+  */
+  set isActive (active) {
+    if (this._isActive !== active) {
+        this._isActive = active
+        this.restartTokenDetection()
+      }
+    }
 }
 
 module.exports = DetectTokensController