fix: add provideComponentStore function for hook logic and update tests

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

@@ -1,8 +1,10 @@
 /* eslint-disable @typescript-eslint/no-non-null-assertion */
 import {
   ComponentStore,
+  INITIAL_STATE_TOKEN,
   OnStateInit,
   OnStoreInit,
+  provideComponentStore,
 } from '@ngrx/component-store';
 import { fakeSchedulers, marbles } from 'rxjs-marbles/jest';
 import {
@@ -28,6 +30,7 @@ import {
   concatMap,
 } from 'rxjs/operators';
 import { createSelector } from '@ngrx/store';
+import { Inject, Injectable, InjectionToken, Injector } from '@angular/core';
 
 describe('Component Store', () => {
   describe('initialization', () => {
@@ -1459,75 +1462,86 @@ describe('Component Store', () => {
 
     const onStoreInitMessage = 'on store init called';
     const onStateInitMessage = 'on state init called';
-    let logs: string[] = [];
+
+    const INIT_STATE = new InjectionToken('Init State');
+
+    @Injectable()
     class LifecycleStore
       extends ComponentStore<LifeCycle>
       implements OnStoreInit, OnStateInit
     {
-      constructor(state?: LifeCycle) {
+      logs: string[] = [];
+      constructor(@Inject(INIT_STATE) state?: LifeCycle) {
         super(state);
       }
 
       logEffect = this.effect(
         tap<void>(() => {
-          logs.push('effect');
+          this.logs.push('effect');
         })
       );
 
       ngrxOnStoreInit() {
-        logs.push(onStoreInitMessage);
+        this.logs.push(onStoreInitMessage);
       }
 
       ngrxOnStateInit() {
-        logs.push(onStateInitMessage);
+        this.logs.push(onStateInitMessage);
       }
     }
 
-    let componentStore: LifecycleStore;
+    function setup(initialState?: LifeCycle) {
+      const injector = Injector.create({
+        providers: [
+          { provide: INIT_STATE, useValue: initialState },
+          provideComponentStore(LifecycleStore),
+        ],
+      });
 
-    beforeEach(() => {
-      logs = [];
-    });
+      return {
+        store: injector.get(LifecycleStore),
+      };
+    }
 
     it('should call the OnInitStore lifecycle hook if defined', () => {
-      componentStore = new LifecycleStore({ init: true });
+      const state = setup({ init: true });
 
-      expect(logs[0]).toBe(onStoreInitMessage);
+      expect(state.store.logs[0]).toBe(onStoreInitMessage);
     });
 
     it('should only call the OnInitStore lifecycle hook once', () => {
-      componentStore = new LifecycleStore({ init: true });
-      expect(logs[0]).toBe(onStoreInitMessage);
+      const state = setup({ init: true });
+      expect(state.store.logs[0]).toBe(onStoreInitMessage);
 
-      logs = [];
-      componentStore.setState({ init: false });
+      state.store.logs = [];
+      state.store.setState({ init: false });
 
-      expect(logs.length).toBe(0);
+      expect(state.store.logs.length).toBe(0);
     });
 
     it('should call the OnInitState lifecycle hook if defined and state is set eagerly', () => {
-      componentStore = new LifecycleStore({ init: true });
+      const state = setup({ init: true });
 
-      expect(logs[1]).toBe(onStateInitMessage);
+      expect(state.store.logs[1]).toBe(onStateInitMessage);
     });
 
     it('should call the OnInitState lifecycle hook if defined and after state is set lazily', () => {
-      componentStore = new LifecycleStore();
-      expect(logs.length).toBe(1);
+      const state = setup();
+      expect(state.store.logs.length).toBe(1);
 
-      componentStore.setState({ init: true });
+      state.store.setState({ init: true });
 
-      expect(logs[1]).toBe(onStateInitMessage);
+      expect(state.store.logs[1]).toBe(onStateInitMessage);
     });
 
     it('should only call the OnInitStore lifecycle hook once', () => {
-      componentStore = new LifecycleStore({ init: true });
+      const state = setup({ init: true });
 
-      expect(logs[1]).toBe(onStateInitMessage);
-      logs = [];
-      componentStore.setState({ init: false });
+      expect(state.store.logs[1]).toBe(onStateInitMessage);
+      state.store.logs = [];
+      state.store.setState({ init: false });
 
-      expect(logs.length).toBe(0);
+      expect(state.store.logs.length).toBe(0);
     });
   });
 });

modules/component-store/src/component-store.ts

@@ -9,7 +9,6 @@ import {
   Subject,
   queueScheduler,
   scheduled,
-  EMPTY,
 } from 'rxjs';
 import {
   concatMap,
@@ -19,7 +18,6 @@ import {
   distinctUntilChanged,
   shareReplay,
   take,
-  catchError,
 } from 'rxjs/operators';
 import { debounceSync } from './debounce-sync';
 import {
@@ -73,19 +71,6 @@ export class ComponentStore<T extends object> implements OnDestroy {
   // Needs to be after destroy$ is declared because it's used in select.
   readonly state$: Observable<T> = this.select((s) => s);
 
-  // // check/call store init hook
-  // private readonly initStoreHook = this.effect(() =>
-  //   of(null).pipe(($) => {
-  //     if (isOnStoreInitDefined(this)) {
-  //       this.ngrxOnStoreInit();
-  //     }
-  //     return $;
-  //   })
-  // )();
-
-  // check/call state init hook on first emission of value
-  // private readonly initStateHook = this.callInitStateHook();
-
   constructor(@Optional() @Inject(INITIAL_STATE_TOKEN) defaultState?: T) {
     // State can be initialized either through constructor or setState.
     if (defaultState) {
@@ -333,24 +318,6 @@ export class ComponentStore<T extends object> implements OnDestroy {
       });
     }) as unknown as ReturnType;
   }
-
-  callInitStateHook() {
-    this.stateSubject$
-      .pipe(
-        take(1),
-        map((val) => {
-          if (val && isOnStateInitDefined(this)) {
-            this.ngrxOnStateInit();
-          }
-          return val;
-        }),
-        catchError((e) => {
-          console.log(e);
-          return EMPTY;
-        })
-      )
-      .subscribe();
-  }
 }
 
 function processSelectorArgs<
@@ -400,27 +367,33 @@ const WITH_HOOKS = new InjectionToken<ComponentStore<any>[]>(
   '@ngrx/component-store: ComponentStores with Hooks'
 );
 
-export function provideWithHooks(
+export function provideComponentStore(
   componentStoreClass: Type<ComponentStore<any>>
 ) {
   return [
     { provide: WITH_HOOKS, multi: true, useClass: componentStoreClass },
     {
       provide: componentStoreClass,
       useFactory: () => {
-        const componentStore = inject(WITH_HOOKS).pop();
-
-        if (isOnStoreInitDefined(componentStore)) {
-          componentStore.ngrxOnStoreInit();
-        }
-
-        if (isOnStateInitDefined(componentStore)) {
-          componentStore.state$
-            .pipe(take(1))
-            .subscribe(() => componentStore.ngrxOnStateInit());
-        }
+        const componentStores = inject(WITH_HOOKS);
+        let instance;
+        componentStores.forEach((componentStore) => {
+          if (componentStore instanceof componentStoreClass) {
+            instance = componentStore;
+
+            if (isOnStoreInitDefined(componentStore)) {
+              componentStore.ngrxOnStoreInit();
+            }
+
+            if (isOnStateInitDefined(componentStore)) {
+              componentStore.state$
+                .pipe(take(1))
+                .subscribe(() => componentStore.ngrxOnStateInit());
+            }
+          }
+        });
 
-        return componentStore;
+        return instance;
       },
     },
   ];

projects/example-app/src/app/auth/containers/login-page.component.ts

@@ -3,8 +3,6 @@ import { Store } from '@ngrx/store';
 import { Credentials } from '@example-app/auth/models';
 import * as fromAuth from '@example-app/auth/reducers';
 import { LoginPageActions } from '@example-app/auth/actions';
-import { LoginPageStore } from './login-page.store';
-import { provideWithHooks } from '@ngrx/component-store';
 
 @Component({
   selector: 'bc-login-page',
@@ -17,17 +15,12 @@ import { provideWithHooks } from '@ngrx/component-store';
     </bc-login-form>
   `,
   styles: [],
-  providers: [provideWithHooks(LoginPageStore)],
 })
 export class LoginPageComponent {
   pending$ = this.store.select(fromAuth.selectLoginPagePending);
   error$ = this.store.select(fromAuth.selectLoginPageError);
 
-  constructor(private store: Store, private lg: LoginPageStore) {}
-
-  ngOnInit() {
-    this.lg.setState({ init: true });
-  }
+  constructor(private store: Store<fromAuth.State>) {}
 
   onSubmit(credentials: Credentials) {
     this.store.dispatch(LoginPageActions.login({ credentials }));