3 files changed, 74 insertions, 3 deletions

diff --git a/app/scripts/keyring-controller.js b/app/scripts/keyring-controller.js
index 5cf2542cc..807752a94 100644
--- a/app/scripts/keyring-controller.js
+++ b/app/scripts/keyring-controller.js
@@ -18,6 +18,8 @@ module.exports = class KeyringController extends EventEmitter {
     this.configManager = opts.configManager
     this.ethStore = opts.ethStore
     this.encryptor = encryptor
+    this.keyringTypes = keyringTypes
+
     this.keyrings = []
     this.identities = {} // Essentially a name hash
   }
@@ -37,7 +39,7 @@ module.exports = class KeyringController extends EventEmitter {
       currentFiat: this.configManager.getCurrentFiat(),
       conversionRate: this.configManager.getConversionRate(),
       conversionDate: this.configManager.getConversionDate(),
-      keyringTypes: keyringTypes.map((krt) => krt.type()),
+      keyringTypes: this.keyringTypes.map((krt) => krt.type()),
       identities: this.identities,
     }
   }
@@ -154,7 +156,7 @@ module.exports = class KeyringController extends EventEmitter {
     })
   }
 
-  restoreKeyring(serialized, i) {
+  restoreKeyring(i, serialized) {
     const { type, data } = serialized
     const Keyring = this.getKeyringClassForType(type)
     const keyring = new Keyring()
@@ -168,7 +170,7 @@ module.exports = class KeyringController extends EventEmitter {
   }
 
   getKeyringClassForType(type) {
-    const Keyring = keyringTypes.reduce((res, kr) => {
+    const Keyring = this.keyringTypes.reduce((res, kr) => {
       if (kr.type() === type) {
         return kr
       } else {
diff --git a/test/lib/mock-simple-keychain.js b/test/lib/mock-simple-keychain.js
new file mode 100644
index 000000000..615b3e537
--- /dev/null
+++ b/test/lib/mock-simple-keychain.js
@@ -0,0 +1,38 @@
+var fakeWallet = {
+  privKey: '0x123456788890abcdef',
+  address: '0xfedcba0987654321',
+}
+const type = 'Simple Key Pair'
+
+module.exports = class MockSimpleKeychain {
+
+  static type() { return type }
+
+  constructor(opts) {
+    this.type = type
+    this.opts = opts || {}
+    this.wallets = []
+  }
+
+  serialize() {
+    return [ fakeWallet.privKey ]
+  }
+
+  deserialize(data) {
+    if (!Array.isArray(data)) {
+      throw new Error('Simple keychain deserialize requires a privKey array.')
+    }
+    this.wallets = [ fakeWallet ]
+  }
+
+  addAccounts(n = 1) {
+    for(var i = 0; i < n; i++) {
+      this.wallets.push(fakeWallet)
+    }
+  }
+
+  getAccounts() {
+    return this.wallets.map(w => w.address)
+  }
+
+}
diff --git a/test/unit/keyring-controller-test.js b/test/unit/keyring-controller-test.js
index 738aa1a84..e216b0960 100644
--- a/test/unit/keyring-controller-test.js
+++ b/test/unit/keyring-controller-test.js
@@ -4,6 +4,8 @@ var configManagerGen = require('../lib/mock-config-manager')
 const ethUtil = require('ethereumjs-util')
 const async = require('async')
 const mockEncryptor = require('../lib/mock-encryptor')
+const MockSimpleKeychain = require('../lib/mock-simple-keychain')
+const sinon = require('sinon')
 
 describe('KeyringController', function() {
 
@@ -15,6 +17,7 @@ describe('KeyringController', function() {
   let originalKeystore
 
   beforeEach(function(done) {
+    this.sinon = sinon.sandbox.create()
     window.localStorage = {} // Hacking localStorage support into JSDom
 
     keyringController = new KeyringController({
@@ -33,8 +36,14 @@ describe('KeyringController', function() {
     })
   })
 
+  afterEach(function() {
+    // Cleanup mocks
+    this.sinon.restore()
+  })
+
   describe('#createNewVault', function () {
     it('should set a vault on the configManager', function(done) {
+      keyringController.configManager.setVault(null)
       assert(!keyringController.configManager.getVault(), 'no previous vault')
       keyringController.createNewVault(password, null, function (err, state) {
         assert.ifError(err)
@@ -44,6 +53,28 @@ describe('KeyringController', function() {
       })
     })
   })
+
+  describe('#restoreKeyring', function(done) {
+
+    it(`should pass a keyring's serialized data back to the correct type.`, function() {
+      keyringController.keyringTypes = [ MockSimpleKeychain ]
+
+      const mockSerialized = {
+        type: MockSimpleKeychain.type(),
+        data: [ '0x123456null788890abcdef' ],
+      }
+      const mock = this.sinon.mock(keyringController)
+
+      mock.expects('loadBalanceAndNickname')
+      .exactly(1)
+
+      var keyring = keyringController.restoreKeyring(0, mockSerialized)
+      assert.equal(keyring.wallets.length, 1, 'one wallet restored')
+      mock.verify()
+    })
+
+  })
+
 })