create tokens storage when setadresses

2 files changed, 29 insertions, 2 deletions

diff --git a/app/scripts/controllers/preferences.js b/app/scripts/controllers/preferences.js
index 4daa196a..5c158288 100644
--- a/app/scripts/controllers/preferences.js
+++ b/app/scripts/controllers/preferences.js
@@ -87,8 +87,10 @@ class PreferencesController {
       ids[address] = {name: `Account ${index + 1}`, address, ...oldId}
       return ids
     }, {})
-    const accountTokens = addresses.reduce((address) => {
-      return oldAccountTokens[address] || {}
+    const accountTokens = addresses.reduce((tokens, address) => {
+      const oldTokens = oldAccountTokens[address] || {}
+      tokens[address] = oldTokens
+      return tokens
     }, {})
     this.store.updateState({ identities, accountTokens })
   }
diff --git a/test/unit/app/controllers/preferences-controller-test.js b/test/unit/app/controllers/preferences-controller-test.js
index dcf56e91..9b2c846b 100644
--- a/test/unit/app/controllers/preferences-controller-test.js
+++ b/test/unit/app/controllers/preferences-controller-test.js
@@ -31,6 +31,20 @@ describe('preferences controller', function () {
       })
     })
 
+    it('should create account tokens for each account in the store', function () {
+      preferencesController.setAddresses([
+        '0xda22le',
+        '0x7e57e2',
+      ])
+
+      const accountTokens = preferencesController.store.getState().accountTokens
+
+      assert.deepEqual(accountTokens, {
+        '0xda22le': {},
+        '0x7e57e2': {},
+      })
+    })
+
     it('should replace its list of addresses', function () {
       preferencesController.setAddresses([
         '0xda22le',
@@ -67,6 +81,17 @@ describe('preferences controller', function () {
       assert.equal(preferencesController.store.getState().identities['0xda22le'], undefined)
     })
 
+    it('should remove an address from state and respective tokens', function () {
+      preferencesController.setAddresses([
+        '0xda22le',
+        '0x7e57e2',
+      ])
+
+      preferencesController.removeAddress('0xda22le')
+
+      assert.equal(preferencesController.store.getState().accountTokens['0xda22le'], undefined)
+    })
+
     it('should switch accounts if the selected address is removed', function () {
       preferencesController.setAddresses([
         '0xda22le',