Fix vault creation bug

2 files changed, 41 insertions, 25 deletions

diff --git a/app/scripts/keyring-controller.js b/app/scripts/keyring-controller.js
index 3b59c7890..3bc9561e2 100644
--- a/app/scripts/keyring-controller.js
+++ b/app/scripts/keyring-controller.js
@@ -78,7 +78,7 @@ module.exports = class KeyringController extends EventEmitter {
   }
 
   createNewVaultAndKeychain(password, entropy, cb) {
-    this.createNewVault(password, entropy, (err, serialized) => {
+    this.createNewVault(password, entropy, (err) => {
       if (err) return cb(err)
       this.createFirstKeyTree(password, cb)
     })
@@ -107,8 +107,8 @@ module.exports = class KeyringController extends EventEmitter {
         const firstAccount = accounts[0]
         const hexAccount = ethUtil.addHexPrefix(firstAccount)
         this.configManager.setSelectedAccount(hexAccount)
-
         this.setupAccounts(accounts)
+
         this.emit('update')
         cb(null, this.getState())
       })
@@ -127,8 +127,7 @@ module.exports = class KeyringController extends EventEmitter {
     })
     .then((serialized) => {
       if (serialized && shouldMigrate) {
-        const accountLength = this.getAccounts().length
-        const keyring = this.restoreKeyring(accountLength, serialized)
+        const keyring = this.restoreKeyring(serialized)
         this.keyrings.push(keyring)
         this.configManager.setSelectedAccount(keyring.getAccounts()[0])
       }
@@ -144,12 +143,8 @@ module.exports = class KeyringController extends EventEmitter {
     configManager.setSalt(salt)
 
     return this.migrateAndGetKey(password)
-    .then((key) => {
-      cb(null, configManager.getVault())
-    })
-    .then((encryptedString) => {
-      const serialized = this.keyrings[0].serialize()
-      cb(null, serialized)
+    .then(() => {
+      cb(null)
     })
     .catch((err) => {
       cb(err)
@@ -160,12 +155,14 @@ module.exports = class KeyringController extends EventEmitter {
     this.clearKeyrings()
     this.addNewKeyring('HD Key Tree', {n: 1}, (err) => {
       const firstKeyring = this.keyrings[0]
-      const firstAccount = firstKeyring.getAccounts()[0]
+      const accounts = firstKeyring.getAccounts()
+      const firstAccount = accounts[0]
       const hexAccount = ethUtil.addHexPrefix(firstAccount)
       const seedWords = firstKeyring.serialize().mnemonic
-      this.configManager.setSelectedAccount(hexAccount)
+      this.configManager.setSelectedAccount(firstAccount)
       this.configManager.setSeedWords(seedWords)
       autoFaucet(hexAccount)
+      this.setupAccounts(accounts)
       this.persistAllKeyrings()
       cb(err, this.getState())
     })
@@ -284,7 +281,7 @@ module.exports = class KeyringController extends EventEmitter {
     const encryptedVault = this.configManager.getVault()
     return this.encryptor.decryptWithKey(key, encryptedVault)
     .then((vault) => {
-      this.keyrings = vault.map(this.restoreKeyring.bind(this, 0))
+      this.keyrings = vault.map(this.restoreKeyring.bind(this))
       return this.persistAllKeyrings()
     })
     .then(() => {
@@ -292,15 +289,14 @@ module.exports = class KeyringController extends EventEmitter {
     })
   }
 
-  restoreKeyring(i, serialized) {
+  restoreKeyring(serialized) {
     const { type, data } = serialized
     const Keyring = this.getKeyringClassForType(type)
     const keyring = new Keyring()
     keyring.deserialize(data)
 
-    keyring.getAccounts().forEach((account) => {
-      this.loadBalanceAndNickname(account, i)
-    })
+    const accounts = keyring.getAccounts()
+    this.setupAccounts(accounts)
 
     this.keyrings.push(keyring)
     return keyring
diff --git a/test/unit/keyring-controller-test.js b/test/unit/keyring-controller-test.js
index 16a4ae148..6656583e8 100644
--- a/test/unit/keyring-controller-test.js
+++ b/test/unit/keyring-controller-test.js
@@ -9,10 +9,11 @@ const sinon = require('sinon')
 
 describe('KeyringController', function() {
 
-  let keyringController
+  let keyringController, state
   let password = 'password123'
   let entropy = 'entripppppyy duuude'
-  let seedWords
+  let seedWords = 'puzzle seed penalty soldier say clay field arctic metal hen cage runway'
+  let addresses = ['eF35cA8EbB9669A35c31b5F6f249A9941a812AC1'.toLowerCase()]
   let accounts = []
   let originalKeystore
 
@@ -31,7 +32,9 @@ describe('KeyringController', function() {
     // Browser crypto is tested in the integration test suite.
     keyringController.encryptor = mockEncryptor
 
-    keyringController.createNewVaultAndKeychain(password, null, function (err, state) {
+    keyringController.createNewVaultAndKeychain(password, null, function (err, newState) {
+      assert.ifError(err)
+      state = newState
       done()
     })
   })
@@ -42,6 +45,8 @@ describe('KeyringController', function() {
   })
 
   describe('#createNewVaultAndKeychain', function () {
+    this.timeout(10000)
+
     it('should set a vault on the configManager', function(done) {
       keyringController.configManager.setVault(null)
       assert(!keyringController.configManager.getVault(), 'no previous vault')
@@ -57,19 +62,21 @@ describe('KeyringController', function() {
   describe('#restoreKeyring', function() {
 
     it(`should pass a keyring's serialized data back to the correct type.`, function() {
-      keyringController.keyringTypes = [ MockSimpleKeychain ]
-
       const mockSerialized = {
-        type: MockSimpleKeychain.type(),
-        data: [ '0x123456null788890abcdef' ],
+        type: 'HD Key Tree',
+        data: {
+          mnemonic: seedWords,
+          n: 1,
+        }
       }
       const mock = this.sinon.mock(keyringController)
 
       mock.expects('loadBalanceAndNickname')
       .exactly(1)
 
-      var keyring = keyringController.restoreKeyring(0, mockSerialized)
+      var keyring = keyringController.restoreKeyring(mockSerialized)
       assert.equal(keyring.wallets.length, 1, 'one wallet restored')
+      assert.equal(keyring.getAccounts()[0], addresses[0])
       mock.verify()
     })
 
@@ -85,6 +92,19 @@ describe('KeyringController', function() {
     })
   })
 
+  describe('#createNickname', function() {
+    it('should add the address to the identities hash', function() {
+      const fakeAddress = '0x12345678'
+      keyringController.createNickname(fakeAddress)
+      const identities = keyringController.identities
+      const identity = identities[fakeAddress]
+      assert.equal(identity.address, fakeAddress)
+
+      const nick = keyringController.configManager.nicknameForWallet(fakeAddress)
+      assert.equal(typeof nick, 'string')
+    })
+  })
+
 })