2 files changed, 20 insertions, 39 deletions
diff --git a/test/unit/idStore-migration-test.js b/test/unit/idStore-migration-test.js
index 47894a458..3aaf4bb94 100644
--- a/test/unit/idStore-migration-test.js
+++ b/test/unit/idStore-migration-test.js
@@ -89,6 +89,9 @@ describe('IdentityStore to KeyringController migration', function() {
         assert(!state.lostAccounts, 'no lost accounts')
         done()
       })
+      .catch((err) => {
+        done(err)
+      })
     })
   })
 })
diff --git a/test/unit/keyring-controller-test.js b/test/unit/keyring-controller-test.js
index d6d2db817..347aa2bdf 100644
--- a/test/unit/keyring-controller-test.js
+++ b/test/unit/keyring-controller-test.js
@@ -1,6 +1,6 @@
-var assert = require('assert')
-var KeyringController = require('../../app/scripts/keyring-controller')
-var configManagerGen = require('../lib/mock-config-manager')
+const assert = require('assert')
+const KeyringController = require('../../app/scripts/keyring-controller')
+const configManagerGen = require('../lib/mock-config-manager')
 const ethUtil = require('ethereumjs-util')
 const BN = ethUtil.BN
 const async = require('async')
@@ -55,17 +55,16 @@ describe('KeyringController', 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')
+      keyringController.store.updateState({ vault: null })
+      assert(!keyringController.store.getState().vault, 'no previous vault')
       keyringController.createNewVaultAndKeychain(password)
       .then(() => {
-        const vault = keyringController.configManager.getVault()
+        const vault = keyringController.store.getState().vault
         assert(vault, 'vault created')
         done()
       })
       .catch((reason) => {
-        assert.ifError(reason)
-        done()
+        done(reason)
       })
     })
   })
@@ -96,8 +95,7 @@ describe('KeyringController', function() {
         done()
       })
       .catch((reason) => {
-        assert.ifError(reason)
-        done()
+        done(reason)
       })
     })
   })
@@ -109,9 +107,6 @@ describe('KeyringController', function() {
       const identities = keyringController.identities
       const identity = identities[fakeAddress]
       assert.equal(identity.address, fakeAddress)
-
-      const nick = keyringController.configManager.nicknameForWallet(fakeAddress)
-      assert.equal(typeof nick, 'string')
     })
   })
 
@@ -122,34 +117,17 @@ describe('KeyringController', function() {
       keyringController.identities[ethUtil.addHexPrefix(account)] = {}
       keyringController.saveAccountLabel(account, nick)
       .then((label) => {
-        assert.equal(label, nick)
-        const persisted = keyringController.configManager.nicknameForWallet(account)
-        assert.equal(persisted, nick)
-        done()
-      })
-      .catch((reason) => {
-        assert.ifError(reason)
-        done()
-      })
-    })
-
-    this.timeout(10000)
-    it('retrieves the persisted nickname', function(done) {
-      const account = addresses[0]
-      var nick = 'Test nickname'
-      keyringController.configManager.setNicknameForWallet(account, nick)
-      keyringController.createNewVaultAndRestore(password, seedWords)
-      .then((state) => {
-
-        const identity = keyringController.identities['0x' + account]
-        assert.equal(identity.name, nick)
-
-        assert(accounts)
-        done()
+        try {
+          assert.equal(label, nick)
+          const persisted = keyringController.store.getState().walletNicknames[account]
+          assert.equal(persisted, nick)
+          done()
+        } catch (err) {
+          done()
+        }
       })
       .catch((reason) => {
-        assert.ifError(reason)
-        done()
+        done(reason)
       })
     })
   })