fix: create injection token inside provideComponentStore

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

@@ -1,7 +1,6 @@
 /* eslint-disable @typescript-eslint/no-non-null-assertion */
 import {
   ComponentStore,
-  INITIAL_STATE_TOKEN,
   OnStateInit,
   OnStoreInit,
   provideComponentStore,
@@ -30,7 +29,13 @@ import {
   concatMap,
 } from 'rxjs/operators';
 import { createSelector } from '@ngrx/store';
-import { Inject, Injectable, InjectionToken, Injector } from '@angular/core';
+import {
+  Inject,
+  Injectable,
+  InjectionToken,
+  Injector,
+  Provider,
+} from '@angular/core';
 
 describe('Component Store', () => {
   describe('initialization', () => {
@@ -1490,27 +1495,39 @@ describe('Component Store', () => {
       }
     }
 
-    function setup(initialState?: LifeCycle) {
+    @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({ init: true });
+      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({ init: true });
+      const state = setup({ initialState: { init: true } });
       expect(state.store.logs[0]).toBe(onStoreInitMessage);
 
       state.store.logs = [];
@@ -1520,7 +1537,7 @@ describe('Component Store', () => {
     });
 
     it('should call the OnInitState lifecycle hook if defined and state is set eagerly', () => {
-      const state = setup({ init: true });
+      const state = setup({ initialState: { init: true } });
 
       expect(state.store.logs[1]).toBe(onStateInitMessage);
     });
@@ -1535,13 +1552,25 @@ describe('Component Store', () => {
     });
 
     it('should only call the OnInitStore lifecycle hook once', () => {
-      const state = setup({ init: true });
+      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/component-store.ts

@@ -26,22 +26,12 @@ import {
   Optional,
   InjectionToken,
   Inject,
-  Type,
-  inject,
 } from '@angular/core';
 
 export interface SelectConfig {
   debounce?: boolean;
 }
 
-export interface OnStoreInit {
-  readonly ngrxOnStoreInit: () => void;
-}
-
-export interface OnStateInit {
-  readonly ngrxOnStateInit: () => void;
-}
-
 export const INITIAL_STATE_TOKEN = new InjectionToken(
   '@ngrx/component-store Initial State'
 );
@@ -354,47 +344,3 @@ function processSelectorArgs<
     config,
   };
 }
-
-function isOnStoreInitDefined(cs: unknown): cs is OnStoreInit {
-  return typeof (cs as OnStoreInit).ngrxOnStoreInit === 'function';
-}
-
-function isOnStateInitDefined(cs: unknown): cs is OnStateInit {
-  return typeof (cs as OnStateInit).ngrxOnStateInit === 'function';
-}
-
-const WITH_HOOKS = new InjectionToken<ComponentStore<any>[]>(
-  '@ngrx/component-store: ComponentStores with Hooks'
-);
-
-export function provideComponentStore(
-  componentStoreClass: Type<ComponentStore<any>>
-) {
-  return [
-    { provide: WITH_HOOKS, multi: true, useClass: componentStoreClass },
-    {
-      provide: componentStoreClass,
-      useFactory: () => {
-        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 instance;
-      },
-    },
-  ];
-}

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,49 @@
+import { Provider, InjectionToken, Type, inject } from '@angular/core';
+import { take } from 'rxjs';
+import { ComponentStore } from './component-store';
+
+export interface OnStoreInit {
+  readonly ngrxOnStoreInit: () => void;
+}
+
+export interface OnStateInit {
+  readonly ngrxOnStateInit: () => void;
+}
+
+export function isOnStoreInitDefined(cs: unknown): cs is OnStoreInit {
+  return typeof (cs as OnStoreInit).ngrxOnStoreInit === 'function';
+}
+
+export function isOnStateInitDefined(cs: unknown): cs is OnStateInit {
+  return typeof (cs as OnStateInit).ngrxOnStateInit === 'function';
+}
+
+export function provideComponentStore(
+  componentStoreClass: Type<ComponentStore<any>>
+): Provider[] {
+  const CS_WITH_HOOKS = new InjectionToken<ComponentStore<any>>(
+    '@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;
+      },
+    },
+  ];
+}