Add PreferencesController#setAddresses to update ids

3 files changed, 66 insertions, 8 deletions

diff --git a/app/scripts/controllers/preferences.js b/app/scripts/controllers/preferences.js
index 55416d15f..a4ff1207e 100644
--- a/app/scripts/controllers/preferences.js
+++ b/app/scripts/controllers/preferences.js
@@ -63,6 +63,16 @@ class PreferencesController {
     this.store.updateState({ currentLocale: key })
   }
 
+  setAddresses (addresses) {
+    const oldIdentities = this.store.getState().identities
+    const identities = addresses.reduce((ids, address, index) => {
+      const oldId = oldIdentities[address] || {}
+      ids[address] = {name: `Account ${index + 1}`, address, ...oldId}
+      return ids
+    }, {})
+    this.store.updateState({ identities })
+  }
+
   /**
    * Setter for the `selectedAddress` property
    *
diff --git a/app/scripts/metamask-controller.js b/app/scripts/metamask-controller.js
index 06ee6c47a..807c9a0b9 100644
--- a/app/scripts/metamask-controller.js
+++ b/app/scripts/metamask-controller.js
@@ -434,7 +434,9 @@ module.exports = class MetamaskController extends EventEmitter {
 
       } else {
         vault = await this.keyringController.createNewVaultAndKeychain(password)
-        this.selectFirstIdentity(vault)
+        const accounts = await this.keyringController.getAccounts()
+        this.preferencesController.setAddresses(accounts)
+        this.selectFirstIdentity()
       }
       release()
     } catch (err) {
@@ -454,7 +456,9 @@ module.exports = class MetamaskController extends EventEmitter {
     const release = await this.createVaultMutex.acquire()
     try {
       const vault = await this.keyringController.createNewVaultAndRestore(password, seed)
-      this.selectFirstIdentity(vault)
+      const accounts = await this.keyringController.getAccounts()
+      this.preferencesController.setAddresses(accounts)
+      this.selectFirstIdentity()
       release()
       return vault
     } catch (err) {
@@ -472,12 +476,10 @@ module.exports = class MetamaskController extends EventEmitter {
    */
 
   /**
-   * Retrieves the first Identiy from the passed Vault and selects the related address
-   *
-   * @param  {} vault
+   * Sets the first address in the state to the selected address
    */
-  selectFirstIdentity (vault) {
-    const { identities } = vault
+  selectFirstIdentity () {
+    const { identities } = this.preferencesController.store.getState()
     const address = Object.keys(identities)[0]
     this.preferencesController.setSelectedAddress(address)
   }
@@ -503,13 +505,15 @@ module.exports = class MetamaskController extends EventEmitter {
 
     await this.verifySeedPhrase()
 
+    this.preferencesController.setAddresses(newAccounts)
     newAccounts.forEach((address) => {
       if (!oldAccounts.includes(address)) {
         this.preferencesController.setSelectedAddress(address)
       }
     })
 
-    return keyState
+    const {identities} = this.preferencesController.store.getState()
+    return {...keyState, identities}
   }
 
   /**
diff --git a/test/unit/preferences-controller-test.js b/test/unit/preferences-controller-test.js
index 8ff1c8648..4b40b1159 100644
--- a/test/unit/preferences-controller-test.js
+++ b/test/unit/preferences-controller-test.js
@@ -8,6 +8,50 @@ describe('preferences controller', function () {
     preferencesController = new PreferencesController()
   })
 
+  describe('setAddresses', function () {
+    it('should keep a map of addresses to names and addresses in the store', function () {
+      preferencesController.setAddresses([
+        '0xda22le',
+        '0x7e57e2',
+      ])
+
+      const {identities} = preferencesController.store.getState()
+      assert.deepEqual(identities, {
+        '0xda22le': {
+          name: 'Account 1',
+          address: '0xda22le',
+        },
+        '0x7e57e2': {
+          name: 'Account 2',
+          address: '0x7e57e2',
+        },
+      })
+    })
+
+    it('should replace its list of addresses', function () {
+      preferencesController.setAddresses([
+        '0xda22le',
+        '0x7e57e2',
+      ])
+      preferencesController.setAddresses([
+        '0xda22le77',
+        '0x7e57e277',
+      ])
+
+      const {identities} = preferencesController.store.getState()
+      assert.deepEqual(identities, {
+        '0xda22le77': {
+          name: 'Account 1',
+          address: '0xda22le77',
+        },
+        '0x7e57e277': {
+          name: 'Account 2',
+          address: '0x7e57e277',
+        },
+      })
+    })
+  })
+
   describe('getTokens', function () {
     it('should return an empty list initially', async function () {
       await preferencesController.setSelectedAddress('0x7e57e2')