4 files changed, 40 insertions, 32 deletions
diff --git a/test/unit/config-manager-test.js b/test/unit/config-manager-test.js
index 77d431d5f..83b242a8b 100644
--- a/test/unit/config-manager-test.js
+++ b/test/unit/config-manager-test.js
@@ -1,5 +1,8 @@
 // polyfill fetch
 global.fetch = global.fetch || require('isomorphic-fetch')
+// pollyfill localStorage support into JSDom
+global.localStorage = global.localStorage || polyfillLocalStorage()
+
 const assert = require('assert')
 const extend = require('xtend')
 const rp = require('request-promise')
@@ -11,7 +14,7 @@ describe('config-manager', function() {
   var configManager
 
   beforeEach(function() {
-    window.localStorage = {} // Hacking localStorage support into JSDom
+    global.localStorage.clear()
     configManager = configManagerGen()
   })
 
@@ -132,7 +135,6 @@ describe('config-manager', function() {
   })
 
   describe('#setConfig', function() {
-    window.localStorage = {} // Hacking localStorage support into JSDom
 
     it('should set the config key', function () {
       var testConfig = {
@@ -238,3 +240,7 @@ describe('config-manager', function() {
     })
   })
 })
+
+function polyfillLocalStorage(){
+  return Object.create({ clear: function(){ global.localStorage = polyfillLocalStorage() } })
+}
diff --git a/test/unit/idStore-migration-test.js b/test/unit/idStore-migration-test.js
index 54f38fb2f..4adbef740 100644
--- a/test/unit/idStore-migration-test.js
+++ b/test/unit/idStore-migration-test.js
@@ -3,6 +3,7 @@ const assert = require('assert')
 const ethUtil = require('ethereumjs-util')
 const BN = ethUtil.BN
 const ConfigManager = require('../../app/scripts/lib/config-manager')
+const ObservableStore = require('../../app/scripts/lib/observable/')
 const delegateCallCode = require('../lib/example-code.json').delegateCallCode
 
 // The old way:
@@ -42,10 +43,9 @@ describe('IdentityStore to KeyringController migration', function() {
   beforeEach(function(done) {
     this.sinon = sinon.sandbox.create()
     window.localStorage = {} // Hacking localStorage support into JSDom
-    configManager = new ConfigManager({
-      loadData,
-      setData: (d) => { window.localStorage = d }
-    })
+    let store = new ObservableStore(loadData())
+    store.subscribe(setData)
+    configManager = new ConfigManager({ store })
 
 
     idStore = new IdentityStore({
diff --git a/test/unit/metamask-controller-test.js b/test/unit/metamask-controller-test.js
index a6164c9a0..24d9ddd67 100644
--- a/test/unit/metamask-controller-test.js
+++ b/test/unit/metamask-controller-test.js
@@ -10,9 +10,11 @@ describe('MetaMaskController', function() {
     showUnconfirmedMessage: noop,
     unlockAccountMessage: noop,
     showUnapprovedTx: noop,
-    setData,
-    loadData,
+    // initial state
+    initState: loadData(),
   })
+  // setup state persistence
+  controller.store.subscribe(setData)
 
   beforeEach(function() {
     // sinon allows stubbing methods that are easily verified
diff --git a/test/unit/migrations-test.js b/test/unit/migrations-test.js
index 9ea8d5c5a..715a5feb0 100644
--- a/test/unit/migrations-test.js
+++ b/test/unit/migrations-test.js
@@ -1,34 +1,34 @@
-var assert = require('assert')
-var path = require('path')
+const assert = require('assert')
+const path = require('path')
 
-var wallet1 = require(path.join('..', 'lib', 'migrations', '001.json'))
+const wallet1 = require(path.join('..', 'lib', 'migrations', '001.json'))
 
-var migration2 = require(path.join('..', '..', 'app', 'scripts', 'migrations', '002'))
-var migration3 = require(path.join('..', '..', 'app', 'scripts', 'migrations', '003'))
-var migration4 = require(path.join('..', '..', 'app', 'scripts', 'migrations', '004'))
+const migration2 = require(path.join('..', '..', 'app', 'scripts', 'migrations', '002'))
+const migration3 = require(path.join('..', '..', 'app', 'scripts', 'migrations', '003'))
+const migration4 = require(path.join('..', '..', 'app', 'scripts', 'migrations', '004'))
 
-describe('wallet1 is migrated successfully', function() {
+const oldTestRpc = 'https://rawtestrpc.metamask.io/'
+const newTestRpc = 'https://testrpc.metamask.io/'
 
-  it('should convert providers', function(done) {
+describe('wallet1 is migrated successfully', function() {
+  it('should convert providers', function() {
 
     wallet1.data.config.provider = { type: 'etherscan', rpcTarget: null }
 
-    var firstResult = migration2.migrate(wallet1.data)
-    assert.equal(firstResult.config.provider.type, 'rpc', 'provider should be rpc')
-    assert.equal(firstResult.config.provider.rpcTarget, 'https://rpc.metamask.io/', 'main provider should be our rpc')
-
-    var oldTestRpc = 'https://rawtestrpc.metamask.io/'
-    var newTestRpc = 'https://testrpc.metamask.io/'
-    firstResult.config.provider.rpcTarget = oldTestRpc
-
-    var secondResult = migration3.migrate(firstResult)
-    assert.equal(secondResult.config.provider.rpcTarget, newTestRpc)
-
-    var thirdResult = migration4.migrate(secondResult)
-    assert.equal(secondResult.config.provider.rpcTarget, null)
-    assert.equal(secondResult.config.provider.type, 'testnet')
-
-    done()
+    return migration2.migrate(wallet1)
+    .then((firstResult) => {
+      assert.equal(firstResult.data.config.provider.type, 'rpc', 'provider should be rpc')
+      assert.equal(firstResult.data.config.provider.rpcTarget, 'https://rpc.metamask.io/', 'main provider should be our rpc')
+      firstResult.data.config.provider.rpcTarget = oldTestRpc
+      return migration3.migrate(firstResult)
+    }).then((secondResult) => {
+      assert.equal(secondResult.data.config.provider.rpcTarget, newTestRpc)
+      return migration4.migrate(secondResult)
+    }).then((thirdResult) => {
+      assert.equal(thirdResult.data.config.provider.rpcTarget, null)
+      assert.equal(thirdResult.data.config.provider.type, 'testnet')
+    })
+    
   })
 })