From 211d7a2626645ebe17dfc711736dcec7438d9c3a Mon Sep 17 00:00:00 2001
From: Matt Gibson <mgibson@bitwarden.com>
Date: Wed, 10 Jan 2024 10:36:19 -0500
Subject: [PATCH] Ps/improve state provider fakers (#7494)

* Expand state provider fakes

- default null initial value for fake states
- Easier mocking of key definitions through just the use of key names
  - allows for not exporting KeyDefinition as long as the key doesn't collide
- mock of fake state provider to verify `get` calls
- `nextMock` for use of the fn mock matchers on emissions of `state$`
- `FakeAccountService` which allows for easy initialization and working with account switching

* Small bug fix for cache key collision on key definitions unique by only storage location

* Fix initial value for test
---
 libs/common/spec/fake-account-service.ts      |  69 +++++++++
 libs/common/spec/fake-state-provider.ts       |  86 +++++++++--
 libs/common/spec/fake-state.ts                | 133 +++++++++++++++---
 .../src/auth/services/account.service.spec.ts |   7 +-
 .../default-state.provider.spec.ts            |   5 +-
 .../src/platform/state/key-definition.ts      |   4 +-
 6 files changed, 269 insertions(+), 35 deletions(-)
 create mode 100644 libs/common/spec/fake-account-service.ts

diff --git a/libs/common/spec/fake-account-service.ts b/libs/common/spec/fake-account-service.ts
new file mode 100644
index 0000000000..f80eb99522
--- /dev/null
+++ b/libs/common/spec/fake-account-service.ts
@@ -0,0 +1,69 @@
+import { mock } from "jest-mock-extended";
+import { Observable, ReplaySubject } from "rxjs";
+
+import { AccountInfo, AccountService } from "../src/auth/abstractions/account.service";
+import { AuthenticationStatus } from "../src/auth/enums/authentication-status";
+import { UserId } from "../src/types/guid";
+
+export function mockAccountServiceWith(
+  userId: UserId,
+  info: Partial<AccountInfo> = {},
+): FakeAccountService {
+  const fullInfo: AccountInfo = {
+    ...info,
+    ...{
+      name: "name",
+      email: "email",
+      status: AuthenticationStatus.Locked,
+    },
+  };
+  const service = new FakeAccountService({ [userId]: fullInfo });
+  service.activeAccountSubject.next({ id: userId, ...fullInfo });
+  return service;
+}
+
+export class FakeAccountService implements AccountService {
+  mock = mock<AccountService>();
+  // eslint-disable-next-line rxjs/no-exposed-subjects -- test class
+  accountsSubject = new ReplaySubject<Record<UserId, AccountInfo>>(1);
+  // eslint-disable-next-line rxjs/no-exposed-subjects -- test class
+  activeAccountSubject = new ReplaySubject<{ id: UserId } & AccountInfo>(1);
+  private _activeUserId: UserId;
+  get activeUserId() {
+    return this._activeUserId;
+  }
+  get accounts$() {
+    return this.accountsSubject.asObservable();
+  }
+  get activeAccount$() {
+    return this.activeAccountSubject.asObservable();
+  }
+  accountLock$: Observable<UserId>;
+  accountLogout$: Observable<UserId>;
+
+  constructor(initialData: Record<UserId, AccountInfo>) {
+    this.accountsSubject.next(initialData);
+    this.activeAccountSubject.subscribe((data) => (this._activeUserId = data?.id));
+    this.activeAccountSubject.next(null);
+  }
+
+  async addAccount(userId: UserId, accountData: AccountInfo): Promise<void> {
+    this.mock.addAccount(userId, accountData);
+  }
+
+  async setAccountName(userId: UserId, name: string): Promise<void> {
+    this.mock.setAccountName(userId, name);
+  }
+
+  async setAccountEmail(userId: UserId, email: string): Promise<void> {
+    this.mock.setAccountEmail(userId, email);
+  }
+
+  async setAccountStatus(userId: UserId, status: AuthenticationStatus): Promise<void> {
+    this.mock.setAccountStatus(userId, status);
+  }
+
+  async switchAccount(userId: UserId): Promise<void> {
+    this.mock.switchAccount(userId);
+  }
+}
diff --git a/libs/common/spec/fake-state-provider.ts b/libs/common/spec/fake-state-provider.ts
index 4712f9077f..5daeb14cb1 100644
--- a/libs/common/spec/fake-state-provider.ts
+++ b/libs/common/spec/fake-state-provider.ts
@@ -1,3 +1,4 @@
+import { mock } from "jest-mock-extended";
 import { Observable } from "rxjs";
 
 import {
@@ -16,6 +17,7 @@ import {
 import { UserId } from "../src/types/guid";
 import { DerivedStateDependencies } from "../src/types/state";
 
+import { FakeAccountService } from "./fake-account-service";
 import {
   FakeActiveUserState,
   FakeDerivedState,
@@ -24,56 +26,114 @@ import {
 } from "./fake-state";
 
 export class FakeGlobalStateProvider implements GlobalStateProvider {
+  mock = mock<GlobalStateProvider>();
+  establishedMocks: Map<string, FakeGlobalState<unknown>> = new Map();
   states: Map<string, GlobalState<unknown>> = new Map();
   get<T>(keyDefinition: KeyDefinition<T>): GlobalState<T> {
-    let result = this.states.get(keyDefinition.buildCacheKey("global")) as GlobalState<T>;
+    this.mock.get(keyDefinition);
+    let result = this.states.get(keyDefinition.buildCacheKey("global"));
 
     if (result == null) {
+      let fake: FakeGlobalState<T>;
+      // Look for established mock
+      if (this.establishedMocks.has(keyDefinition.key)) {
+        fake = this.establishedMocks.get(keyDefinition.key) as FakeGlobalState<T>;
+      } else {
+        fake = new FakeGlobalState<T>();
+      }
+      fake.keyDefinition = keyDefinition;
+      result = fake;
+      this.states.set(keyDefinition.buildCacheKey("global"), result);
+
       result = new FakeGlobalState<T>();
       this.states.set(keyDefinition.buildCacheKey("global"), result);
     }
-    return result;
+    return result as GlobalState<T>;
   }
 
   getFake<T>(keyDefinition: KeyDefinition<T>): FakeGlobalState<T> {
     return this.get(keyDefinition) as FakeGlobalState<T>;
   }
+
+  mockFor<T>(keyDefinitionKey: string, initialValue?: T): FakeGlobalState<T> {
+    if (!this.establishedMocks.has(keyDefinitionKey)) {
+      this.establishedMocks.set(keyDefinitionKey, new FakeGlobalState<T>(initialValue));
+    }
+    return this.establishedMocks.get(keyDefinitionKey) as FakeGlobalState<T>;
+  }
 }
 
 export class FakeSingleUserStateProvider implements SingleUserStateProvider {
+  mock = mock<SingleUserStateProvider>();
+  establishedMocks: Map<string, FakeSingleUserState<unknown>> = new Map();
   states: Map<string, SingleUserState<unknown>> = new Map();
   get<T>(userId: UserId, keyDefinition: KeyDefinition<T>): SingleUserState<T> {
-    let result = this.states.get(keyDefinition.buildCacheKey("user", userId)) as SingleUserState<T>;
+    this.mock.get(userId, keyDefinition);
+    let result = this.states.get(keyDefinition.buildCacheKey("user", userId));
 
     if (result == null) {
-      result = new FakeSingleUserState<T>(userId);
+      let fake: FakeSingleUserState<T>;
+      // Look for established mock
+      if (this.establishedMocks.has(keyDefinition.key)) {
+        fake = this.establishedMocks.get(keyDefinition.key) as FakeSingleUserState<T>;
+      } else {
+        fake = new FakeSingleUserState<T>(userId);
+      }
+      fake.keyDefinition = keyDefinition;
+      result = fake;
       this.states.set(keyDefinition.buildCacheKey("user", userId), result);
     }
-    return result;
+    return result as SingleUserState<T>;
   }
 
   getFake<T>(userId: UserId, keyDefinition: KeyDefinition<T>): FakeSingleUserState<T> {
     return this.get(userId, keyDefinition) as FakeSingleUserState<T>;
   }
+
+  mockFor<T>(userId: UserId, keyDefinitionKey: string, initialValue?: T): FakeSingleUserState<T> {
+    if (!this.establishedMocks.has(keyDefinitionKey)) {
+      this.establishedMocks.set(keyDefinitionKey, new FakeSingleUserState<T>(userId, initialValue));
+    }
+    return this.establishedMocks.get(keyDefinitionKey) as FakeSingleUserState<T>;
+  }
 }
 
 export class FakeActiveUserStateProvider implements ActiveUserStateProvider {
-  states: Map<string, ActiveUserState<unknown>> = new Map();
+  establishedMocks: Map<string, FakeActiveUserState<unknown>> = new Map();
+
+  states: Map<string, FakeActiveUserState<unknown>> = new Map();
+
+  constructor(public accountService: FakeAccountService) {}
+
   get<T>(keyDefinition: KeyDefinition<T>): ActiveUserState<T> {
-    let result = this.states.get(
-      keyDefinition.buildCacheKey("user", "active"),
-    ) as ActiveUserState<T>;
+    let result = this.states.get(keyDefinition.buildCacheKey("user", "active"));
 
     if (result == null) {
-      result = new FakeActiveUserState<T>();
+      // Look for established mock
+      if (this.establishedMocks.has(keyDefinition.key)) {
+        result = this.establishedMocks.get(keyDefinition.key);
+      } else {
+        result = new FakeActiveUserState<T>(this.accountService);
+      }
+      result.keyDefinition = keyDefinition;
       this.states.set(keyDefinition.buildCacheKey("user", "active"), result);
     }
-    return result;
+    return result as ActiveUserState<T>;
   }
 
   getFake<T>(keyDefinition: KeyDefinition<T>): FakeActiveUserState<T> {
     return this.get(keyDefinition) as FakeActiveUserState<T>;
   }
+
+  mockFor<T>(keyDefinitionKey: string, initialValue?: T): FakeActiveUserState<T> {
+    if (!this.establishedMocks.has(keyDefinitionKey)) {
+      this.establishedMocks.set(
+        keyDefinitionKey,
+        new FakeActiveUserState<T>(this.accountService, initialValue),
+      );
+    }
+    return this.establishedMocks.get(keyDefinitionKey) as FakeActiveUserState<T>;
+  }
 }
 
 export class FakeStateProvider implements StateProvider {
@@ -97,9 +157,11 @@ export class FakeStateProvider implements StateProvider {
     return this.derived.get(parentState$, deriveDefinition, dependencies);
   }
 
+  constructor(public accountService: FakeAccountService) {}
+
   global: FakeGlobalStateProvider = new FakeGlobalStateProvider();
   singleUser: FakeSingleUserStateProvider = new FakeSingleUserStateProvider();
-  activeUser: FakeActiveUserStateProvider = new FakeActiveUserStateProvider();
+  activeUser: FakeActiveUserStateProvider = new FakeActiveUserStateProvider(this.accountService);
   derived: FakeDerivedStateProvider = new FakeDerivedStateProvider();
 }
 
diff --git a/libs/common/spec/fake-state.ts b/libs/common/spec/fake-state.ts
index b56c129c62..6c29c20ad8 100644
--- a/libs/common/spec/fake-state.ts
+++ b/libs/common/spec/fake-state.ts
@@ -1,12 +1,20 @@
 import { Observable, ReplaySubject, firstValueFrom, map, timeout } from "rxjs";
 
-import { DerivedState, GlobalState, SingleUserState, ActiveUserState } from "../src/platform/state";
+import {
+  DerivedState,
+  GlobalState,
+  SingleUserState,
+  ActiveUserState,
+  KeyDefinition,
+} from "../src/platform/state";
 // eslint-disable-next-line import/no-restricted-paths -- using unexposed options for clean typing in test class
 import { StateUpdateOptions } from "../src/platform/state/state-update-options";
 // eslint-disable-next-line import/no-restricted-paths -- using unexposed options for clean typing in test class
-import { CombinedState, UserState, activeMarker } from "../src/platform/state/user-state";
+import { CombinedState, activeMarker } from "../src/platform/state/user-state";
 import { UserId } from "../src/types/guid";
 
+import { FakeAccountService } from "./fake-account-service";
+
 const DEFAULT_TEST_OPTIONS: StateUpdateOptions<any, any> = {
   shouldUpdate: () => true,
   combineLatestWith: null,
@@ -26,6 +34,10 @@ export class FakeGlobalState<T> implements GlobalState<T> {
   // eslint-disable-next-line rxjs/no-exposed-subjects -- exposed for testing setup
   stateSubject = new ReplaySubject<T>(1);
 
+  constructor(initialValue?: T) {
+    this.stateSubject.next(initialValue ?? null);
+  }
+
   update: <TCombine>(
     configureState: (state: T, dependency: TCombine) => T,
     options?: StateUpdateOptions<T, TCombine>,
@@ -47,34 +59,56 @@ export class FakeGlobalState<T> implements GlobalState<T> {
     }
     const newState = configureState(current, combinedDependencies);
     this.stateSubject.next(newState);
+    this.nextMock(newState);
     return newState;
   });
 
   updateMock = this.update as jest.MockedFunction<typeof this.update>;
+  nextMock = jest.fn<void, [T]>();
 
   get state$() {
     return this.stateSubject.asObservable();
   }
+
+  private _keyDefinition: KeyDefinition<T> | null = null;
+  get keyDefinition() {
+    if (this._keyDefinition == null) {
+      throw new Error(
+        "Key definition not yet set, usually this means your sut has not asked for this state yet",
+      );
+    }
+    return this._keyDefinition;
+  }
+  set keyDefinition(value: KeyDefinition<T>) {
+    this._keyDefinition = value;
+  }
 }
 
-abstract class FakeUserState<T> implements UserState<T> {
+export class FakeSingleUserState<T> implements SingleUserState<T> {
   // eslint-disable-next-line rxjs/no-exposed-subjects -- exposed for testing setup
   stateSubject = new ReplaySubject<CombinedState<T>>(1);
 
-  protected userId: UserId;
-
   state$: Observable<T>;
   combinedState$: Observable<CombinedState<T>>;
 
-  constructor() {
+  constructor(
+    readonly userId: UserId,
+    initialValue?: T,
+  ) {
+    this.stateSubject.next([userId, initialValue ?? null]);
+
     this.combinedState$ = this.stateSubject.asObservable();
     this.state$ = this.combinedState$.pipe(map(([_userId, state]) => state));
   }
 
-  update: <TCombine>(
+  nextState(state: T) {
+    this.stateSubject.next([this.userId, state]);
+  }
+
+  async update<TCombine>(
     configureState: (state: T, dependency: TCombine) => T,
     options?: StateUpdateOptions<T, TCombine>,
-  ) => Promise<T> = jest.fn(async (configureState, options) => {
+  ): Promise<T> {
     options = populateOptionsWithDefault(options);
     const current = await firstValueFrom(this.state$.pipe(timeout(options.msTimeout)));
     const combinedDependencies =
@@ -86,22 +120,87 @@ abstract class FakeUserState<T> implements UserState<T> {
     }
     const newState = configureState(current, combinedDependencies);
     this.stateSubject.next([this.userId, newState]);
+    this.nextMock(newState);
     return newState;
-  });
+  }
 
   updateMock = this.update as jest.MockedFunction<typeof this.update>;
-}
 
-export class FakeSingleUserState<T> extends FakeUserState<T> implements SingleUserState<T> {
-  constructor(readonly userId: UserId) {
-    super();
-    this.userId = userId;
+  nextMock = jest.fn<void, [T]>();
+  private _keyDefinition: KeyDefinition<T> | null = null;
+  get keyDefinition() {
+    if (this._keyDefinition == null) {
+      throw new Error(
+        "Key definition not yet set, usually this means your sut has not asked for this state yet",
+      );
+    }
+    return this._keyDefinition;
+  }
+  set keyDefinition(value: KeyDefinition<T>) {
+    this._keyDefinition = value;
   }
 }
-export class FakeActiveUserState<T> extends FakeUserState<T> implements ActiveUserState<T> {
+export class FakeActiveUserState<T> implements ActiveUserState<T> {
   [activeMarker]: true;
-  changeActiveUser(userId: UserId) {
-    this.userId = userId;
+
+  // eslint-disable-next-line rxjs/no-exposed-subjects -- exposed for testing setup
+  stateSubject = new ReplaySubject<CombinedState<T>>(1);
+
+  state$: Observable<T>;
+  combinedState$: Observable<CombinedState<T>>;
+
+  constructor(
+    private accountService: FakeAccountService,
+    initialValue?: T,
+  ) {
+    this.stateSubject.next([accountService.activeUserId, initialValue ?? null]);
+
+    this.combinedState$ = this.stateSubject.asObservable();
+    this.state$ = this.combinedState$.pipe(map(([_userId, state]) => state));
+  }
+
+  get userId() {
+    return this.accountService.activeUserId;
+  }
+
+  nextState(state: T) {
+    this.stateSubject.next([this.userId, state]);
+  }
+
+  async update<TCombine>(
+    configureState: (state: T, dependency: TCombine) => T,
+    options?: StateUpdateOptions<T, TCombine>,
+  ): Promise<T> {
+    options = populateOptionsWithDefault(options);
+    const current = await firstValueFrom(this.state$.pipe(timeout(options.msTimeout)));
+    const combinedDependencies =
+      options.combineLatestWith != null
+        ? await firstValueFrom(options.combineLatestWith.pipe(timeout(options.msTimeout)))
+        : null;
+    if (!options.shouldUpdate(current, combinedDependencies)) {
+      return current;
+    }
+    const newState = configureState(current, combinedDependencies);
+    this.stateSubject.next([this.userId, newState]);
+    this.nextMock(this.userId, newState);
+    return newState;
+  }
+
+  updateMock = this.update as jest.MockedFunction<typeof this.update>;
+
+  nextMock = jest.fn<void, [UserId, T]>();
+
+  private _keyDefinition: KeyDefinition<T> | null = null;
+  get keyDefinition() {
+    if (this._keyDefinition == null) {
+      throw new Error(
+        "Key definition not yet set, usually this means your sut has not asked for this state yet",
+      );
+    }
+    return this._keyDefinition;
+  }
+  set keyDefinition(value: KeyDefinition<T>) {
+    this._keyDefinition = value;
   }
 }
 
diff --git a/libs/common/src/auth/services/account.service.spec.ts b/libs/common/src/auth/services/account.service.spec.ts
index 0f611f8345..14ff8724ff 100644
--- a/libs/common/src/auth/services/account.service.spec.ts
+++ b/libs/common/src/auth/services/account.service.spec.ts
@@ -36,8 +36,6 @@ describe("accountService", () => {
     sut = new AccountServiceImplementation(messagingService, logService, globalStateProvider);
 
     accountsState = globalStateProvider.getFake(ACCOUNT_ACCOUNTS);
-    // initialize to empty
-    accountsState.stateSubject.next({});
     activeAccountIdState = globalStateProvider.getFake(ACCOUNT_ACTIVE_ACCOUNT_ID);
   });
 
@@ -57,7 +55,10 @@ describe("accountService", () => {
       accountsState.stateSubject.next({ [userId]: userInfo(AuthenticationStatus.Unlocked) });
       activeAccountIdState.stateSubject.next(userId);
 
-      expect(emissions).toEqual([{ id: userId, ...userInfo(AuthenticationStatus.Unlocked) }]);
+      expect(emissions).toEqual([
+        undefined, // initial value
+        { id: userId, ...userInfo(AuthenticationStatus.Unlocked) },
+      ]);
     });
 
     it("should update the status if the account status changes", async () => {
diff --git a/libs/common/src/platform/state/implementations/default-state.provider.spec.ts b/libs/common/src/platform/state/implementations/default-state.provider.spec.ts
index 4c86a1b8fd..704dbbccc5 100644
--- a/libs/common/src/platform/state/implementations/default-state.provider.spec.ts
+++ b/libs/common/src/platform/state/implementations/default-state.provider.spec.ts
@@ -1,5 +1,6 @@
 import { of } from "rxjs";
 
+import { FakeAccountService, mockAccountServiceWith } from "../../../../spec/fake-account-service";
 import {
   FakeActiveUserStateProvider,
   FakeDerivedStateProvider,
@@ -19,9 +20,11 @@ describe("DefaultStateProvider", () => {
   let singleUserStateProvider: FakeSingleUserStateProvider;
   let globalStateProvider: FakeGlobalStateProvider;
   let derivedStateProvider: FakeDerivedStateProvider;
+  let accountService: FakeAccountService;
 
   beforeEach(() => {
-    activeUserStateProvider = new FakeActiveUserStateProvider();
+    accountService = mockAccountServiceWith("fakeUserId" as UserId);
+    activeUserStateProvider = new FakeActiveUserStateProvider(accountService);
     singleUserStateProvider = new FakeSingleUserStateProvider();
     globalStateProvider = new FakeGlobalStateProvider();
     derivedStateProvider = new FakeDerivedStateProvider();
diff --git a/libs/common/src/platform/state/key-definition.ts b/libs/common/src/platform/state/key-definition.ts
index 47b83a4e88..4f2481b33f 100644
--- a/libs/common/src/platform/state/key-definition.ts
+++ b/libs/common/src/platform/state/key-definition.ts
@@ -151,8 +151,8 @@ export class KeyDefinition<T> {
       throw new Error("You must provide a userId when building a user scoped cache key.");
     }
     return userId === null
-      ? `${scope}_${userId}_${this.stateDefinition.name}_${this.key}`
-      : `${scope}_${this.stateDefinition.name}_${this.key}`;
+      ? `${this.stateDefinition.storageLocation}_${scope}_${userId}_${this.stateDefinition.name}_${this.key}`
+      : `${this.stateDefinition.storageLocation}_${scope}_${this.stateDefinition.name}_${this.key}`;
   }
 
   private get errorKeyName() {