detect tokens when submit password and new account selected

2 files changed, 43 insertions, 17 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
diff --git a/test/unit/app/controllers/detect-tokens-test.js b/test/unit/app/controllers/detect-tokens-test.js
index 860ed7050..49492c543 100644
--- a/test/unit/app/controllers/detect-tokens-test.js
+++ b/test/unit/app/controllers/detect-tokens-test.js
@@ -28,7 +28,7 @@ describe('DetectTokensController', () => {
     const controller = new DetectTokensController({preferences: preferences, network: network})
     controller.isActive = true
 
-    var stub = sandbox.stub(controller, 'exploreNewTokens')
+    var stub = sandbox.stub(controller, 'detectNewTokens')
 
     clock.tick(1)
     sandbox.assert.notCalled(stub)
@@ -51,7 +51,7 @@ describe('DetectTokensController', () => {
         .withArgs('0x0D262e5dC4A06a0F1c90cE79C7a60C09DfC884E4').returns(true)
         .withArgs('0xBC86727E770de68B1060C91f6BB6945c73e10388').returns(true)
 
-    await controller.exploreNewTokens()
+    await controller.detectNewTokens()
     sandbox.assert.notCalled(stub)
   })
 
@@ -68,7 +68,7 @@ describe('DetectTokensController', () => {
         .withArgs('0xBC86727E770de68B1060C91f6BB6945c73e10388')
         .returns(preferences.addToken('0xbc86727e770de68b1060c91f6bb6945c73e10388', 'XNK', 18))
 
-    await controller.exploreNewTokens()
+    await controller.detectNewTokens()
     assert.deepEqual(preferences.store.getState().tokens, [{address: '0x0d262e5dc4a06a0f1c90ce79c7a60c09dfc884e4', decimals: 8, symbol: 'J8T'},
         {address: '0xbc86727e770de68b1060c91f6bb6945c73e10388', decimals: 18, symbol: 'XNK'}])
   })
@@ -87,7 +87,7 @@ describe('DetectTokensController', () => {
       .withArgs('0xBC86727E770de68B1060C91f6BB6945c73e10388')
       .returns(preferences.addToken('0xbc86727e770de68b1060c91f6bb6945c73e10388', 'XNK', 18))
 
-    await controller.exploreNewTokens()
+    await controller.detectNewTokens()
     assert.deepEqual(preferences.store.getState().tokens, [{address: '0x0d262e5dc4a06a0f1c90ce79c7a60c09dfc884e4', decimals: 8, symbol: 'J8T'},
         {address: '0xbc86727e770de68b1060c91f6bb6945c73e10388', decimals: 18, symbol: 'XNK'}])
   })