Update selectedAddress when identities are updated (#415)

The `selectedAddress` state is now reset to the first identity if the
current selected address is removed during an update. This is to ensure
the `selectedAddress` is never set to an address that is not present in
state.

This was already being done manually after each invocation of
`updateIdentities` where account removal was anticipated. For mobile
this is mostly a non-functional change, except that now this update
happens in a single state update rather than over the course of two
successive updates.

src/keyring/KeyringController.ts

@@ -197,7 +197,6 @@ export class KeyringController extends BaseController<KeyringConfig, KeyringStat
       preferences.updateIdentities([]);
       const vault = await privates.get(this).keyring.createNewVaultAndRestore(password, seed);
       preferences.updateIdentities(await privates.get(this).keyring.getAccounts());
-      preferences.update({ selectedAddress: Object.keys(preferences.state.identities)[0] });
       this.fullUpdate();
       return vault;
     } finally {
@@ -217,7 +216,6 @@ export class KeyringController extends BaseController<KeyringConfig, KeyringStat
     try {
       const vault = await privates.get(this).keyring.createNewVaultAndKeychain(password);
       preferences.updateIdentities(await privates.get(this).keyring.getAccounts());
-      preferences.update({ selectedAddress: Object.keys(preferences.state.identities)[0] });
       this.fullUpdate();
       return vault;
     } finally {

src/user/PreferencesController.test.ts

@@ -60,7 +60,7 @@ describe('PreferencesController', () => {
     expect(controller.state.selectedAddress).toBe('0xfoO');
   });
 
-  it('should update existing identities', () => {
+  it('should add new identities', () => {
     const controller = new PreferencesController();
     controller.updateIdentities(['foo', 'bar']);
     expect(controller.state.identities).toEqual({
@@ -69,6 +69,65 @@ describe('PreferencesController', () => {
     });
   });
 
+  it('should not update existing identities', () => {
+    const controller = new PreferencesController(
+      {},
+      { identities: { '0xbar': { address: '0xbar', name: 'Custom name' } } },
+    );
+    controller.updateIdentities(['foo', 'bar']);
+    expect(controller.state.identities).toEqual({
+      '0xbar': { address: '0xbar', name: 'Custom name' },
+      '0xfoO': { address: '0xfoO', name: 'Account 1' },
+    });
+  });
+
+  it('should remove identities', () => {
+    const controller = new PreferencesController(
+      {},
+      {
+        identities: {
+          '0xbar': { address: '0xbar', name: 'Account 2' },
+          '0xfoO': { address: '0xfoO', name: 'Account 1' },
+        },
+      },
+    );
+    controller.updateIdentities(['foo']);
+    expect(controller.state.identities).toEqual({
+      '0xfoO': { address: '0xfoO', name: 'Account 1' },
+    });
+  });
+
+  it('should not update selected address if it is still among identities', () => {
+    const controller = new PreferencesController(
+      {},
+      {
+        identities: {
+          '0xbar': { address: '0xbar', name: 'Account 2' },
+          '0xfoO': { address: '0xfoO', name: 'Account 1' },
+        },
+        selectedAddress: '0xbar',
+      },
+    );
+    controller.updateIdentities(['foo', 'bar']);
+    expect(controller.state.selectedAddress).toEqual('0xbar');
+  });
+
+  it('should update selected address to first identity if it was removed from identities', () => {
+    const controller = new PreferencesController(
+      {},
+      {
+        identities: {
+          '0xbar': { address: '0xbar', name: 'Account 2' },
+          '0xbaz': { address: '0xbaz', name: 'Account 3' },
+          '0xfoO': { address: '0xfoO', name: 'Account 1' },
+        },
+        selectedAddress: '0xbaz',
+      },
+    );
+    controller.updateIdentities(['foo', 'bar']);
+    expect(controller.state.selectedAddress).toEqual('0xfoO');
+  });
+
   it('should add custom rpc url', () => {
     const controller = new PreferencesController();
     const rpcUrlNetwork = {

src/user/PreferencesController.ts

@@ -174,7 +174,9 @@ export class PreferencesController extends BaseController<BaseConfig, Preference
   }
 
   /**
-   * Generates and stores a new list of stored identities based on address
+   * Generates and stores a new list of stored identities based on address. If the selected address
+   * is unset, or if it refers to an identity that was removed, it will be set to the first
+   * identity.
    *
    * @param addresses - List of addresses to use as a basis for each identity
    */
@@ -188,7 +190,11 @@ export class PreferencesController extends BaseController<BaseConfig, Preference
       };
       return ids;
     }, {});
-    this.update({ identities: { ...identities } });
+    let { selectedAddress } = this.state;
+    if (!Object.keys(identities).includes(selectedAddress)) {
+      selectedAddress = Object.keys(identities)[0];
+    }
+    this.update({ identities: { ...identities }, selectedAddress });
   }
 
   /**