feat(component-store): add OnStoreInit and OnStateInit lifecycle hooks (#3368)

Closes #3335

Author: brandonroberts

modules/component-store/spec/component-store.spec.ts

@@ -1,5 +1,10 @@
 /* eslint-disable @typescript-eslint/no-non-null-assertion */
-import { ComponentStore } from '@ngrx/component-store';
+import {
+  ComponentStore,
+  OnStateInit,
+  OnStoreInit,
+  provideComponentStore,
+} from '@ngrx/component-store';
 import { fakeSchedulers, marbles } from 'rxjs-marbles/jest';
 import {
   of,
@@ -24,6 +29,13 @@ import {
   concatMap,
 } from 'rxjs/operators';
 import { createSelector } from '@ngrx/store';
+import {
+  Inject,
+  Injectable,
+  InjectionToken,
+  Injector,
+  Provider,
+} from '@angular/core';
 
 describe('Component Store', () => {
   describe('initialization', () => {
@@ -1447,4 +1459,118 @@ describe('Component Store', () => {
       expect(componentStore.get()).toEqual({ value: 'updated' });
     });
   });
+
+  describe('lifecycle hooks', () => {
+    interface LifeCycle {
+      init: boolean;
+    }
+
+    const onStoreInitMessage = 'on store init called';
+    const onStateInitMessage = 'on state init called';
+
+    const INIT_STATE = new InjectionToken('Init State');
+
+    @Injectable()
+    class LifecycleStore
+      extends ComponentStore<LifeCycle>
+      implements OnStoreInit, OnStateInit
+    {
+      logs: string[] = [];
+      constructor(@Inject(INIT_STATE) state?: LifeCycle) {
+        super(state);
+      }
+
+      logEffect = this.effect(
+        tap<void>(() => {
+          this.logs.push('effect');
+        })
+      );
+
+      ngrxOnStoreInit() {
+        this.logs.push(onStoreInitMessage);
+      }
+
+      ngrxOnStateInit() {
+        this.logs.push(onStateInitMessage);
+      }
+    }
+
+    @Injectable()
+    class ExtraStore extends LifecycleStore {
+      constructor() {
+        super();
+      }
+    }
+
+    function setup({
+      initialState,
+      providers = [],
+    }: { initialState?: LifeCycle; providers?: Provider[] } = {}) {
+      const injector = Injector.create({
+        providers: [
+          { provide: INIT_STATE, useValue: initialState },
+          provideComponentStore(LifecycleStore),
+          providers,
+        ],
+      });
+
+      return {
+        store: injector.get(LifecycleStore),
+        injector,
+      };
+    }
+
+    it('should call the OnInitStore lifecycle hook if defined', () => {
+      const state = setup({ initialState: { init: true } });
+
+      expect(state.store.logs[0]).toBe(onStoreInitMessage);
+    });
+
+    it('should only call the OnInitStore lifecycle hook once', () => {
+      const state = setup({ initialState: { init: true } });
+      expect(state.store.logs[0]).toBe(onStoreInitMessage);
+
+      state.store.logs = [];
+      state.store.setState({ init: false });
+
+      expect(state.store.logs.length).toBe(0);
+    });
+
+    it('should call the OnInitState lifecycle hook if defined and state is set eagerly', () => {
+      const state = setup({ initialState: { init: true } });
+
+      expect(state.store.logs[1]).toBe(onStateInitMessage);
+    });
+
+    it('should call the OnInitState lifecycle hook if defined and after state is set lazily', () => {
+      const state = setup();
+      expect(state.store.logs.length).toBe(1);
+
+      state.store.setState({ init: true });
+
+      expect(state.store.logs[1]).toBe(onStateInitMessage);
+    });
+
+    it('should only call the OnInitStore lifecycle hook once', () => {
+      const state = setup({ initialState: { init: true } });
+
+      expect(state.store.logs[1]).toBe(onStateInitMessage);
+      state.store.logs = [];
+      state.store.setState({ init: false });
+
+      expect(state.store.logs.length).toBe(0);
+    });
+
+    it('works with multiple stores where one extends the other', () => {
+      const state = setup({
+        providers: [provideComponentStore(ExtraStore)],
+      });
+
+      const lifecycleStore = state.store;
+      const extraStore = state.injector.get(ExtraStore);
+
+      expect(lifecycleStore).toBeDefined();
+      expect(extraStore).toBeDefined();
+    });
+  });
 });

modules/component-store/src/index.ts

@@ -1,2 +1,3 @@
 export * from './component-store';
 export * from './tap-response';
+export * from './lifecycle_hooks';

modules/component-store/src/lifecycle_hooks.ts

@@ -0,0 +1,116 @@
+import { Provider, InjectionToken, Type, inject } from '@angular/core';
+import { take } from 'rxjs';
+import { ComponentStore } from './component-store';
+
+/**
+ * The interface for the lifecycle hook
+ * called after the ComponentStore is instantiated.
+ */
+export interface OnStoreInit {
+  readonly ngrxOnStoreInit: () => void;
+}
+
+/**
+ * The interface for the lifecycle hook
+ * called only once after the ComponentStore
+ * state is first initialized.
+ */
+export interface OnStateInit {
+  readonly ngrxOnStateInit: () => void;
+}
+
+/**
+ * Checks to see if the OnInitStore lifecycle hook
+ * is defined on the ComponentStore.
+ *
+ * @param cs ComponentStore type
+ * @returns boolean
+ */
+function isOnStoreInitDefined(cs: unknown): cs is OnStoreInit {
+  return typeof (cs as OnStoreInit).ngrxOnStoreInit === 'function';
+}
+
+/**
+ * Checks to see if the OnInitState lifecycle hook
+ * is defined on the ComponentStore.
+ *
+ * @param cs ComponentStore type
+ * @returns boolean
+ */
+function isOnStateInitDefined(cs: unknown): cs is OnStateInit {
+  return typeof (cs as OnStateInit).ngrxOnStateInit === 'function';
+}
+
+/**
+ * @description
+ *
+ * Function that returns the ComponentStore
+ * class registered as a provider,
+ * and uses a factory provider to instantiate the
+ * ComponentStore and run the lifecycle hooks
+ * defined on the ComponentStore.
+ *
+ * @param componentStoreClass The ComponentStore with lifecycle hooks
+ * @returns Provider[]
+ *
+ * @usageNotes
+ *
+ * ```ts
+ * @Injectable()
+ * export class MyStore
+ *    extends ComponentStore<{ init: boolean }>
+ *    implements OnStoreInit, OnStateInit
+ *   {
+ *
+ *   constructor() {
+ *     super({ init: true });
+ *   }
+ *
+ *   ngrxOnStoreInit() {
+ *     // runs once after store has been instantiated
+ *   }
+ *
+ *   ngrxOnStateInit() {
+ *     // runs once after store state has been initialized
+ *   }
+ * }
+ *
+ * @Component({
+ *   providers: [
+ *     provideComponentStore(MyStore)
+ *   ]
+ * })
+ * export class MyComponent {
+ *   constructor(private myStore: MyStore) {}
+ * }
+ * ```
+ */
+export function provideComponentStore<T extends object>(
+  componentStoreClass: Type<ComponentStore<T>>
+): Provider[] {
+  const CS_WITH_HOOKS = new InjectionToken<ComponentStore<T>>(
+    '@ngrx/component-store ComponentStore with Hooks'
+  );
+
+  return [
+    { provide: CS_WITH_HOOKS, useClass: componentStoreClass },
+    {
+      provide: componentStoreClass,
+      useFactory: () => {
+        const componentStore = inject(CS_WITH_HOOKS);
+
+        if (isOnStoreInitDefined(componentStore)) {
+          componentStore.ngrxOnStoreInit();
+        }
+
+        if (isOnStateInitDefined(componentStore)) {
+          componentStore.state$
+            .pipe(take(1))
+            .subscribe(() => componentStore.ngrxOnStateInit());
+        }
+
+        return componentStore;
+      },
+    },
+  ];
+}