feat(signals): add withHooks signature with factory input (#4208)

Closes #4201

Author: rainerhahnekamp

modules/signals/spec/signal-store.spec.ts

@@ -268,32 +268,58 @@ describe('signalStore', () => {
       expect(message).toBe('onDestroy');
     });
 
-    // FIX: injection context will be provided for `onDestroy` in a separate PR
-    // see https://github.com/ngrx/platform/pull/4196#issuecomment-1875228588
-    it('executes hooks in injection context', () => {
+    it('executes hooks factory in injection context', () => {
       const messages: string[] = [];
-      const TOKEN = new InjectionToken('TOKEN', {
+      const TOKEN_INIT = new InjectionToken('TOKEN_INIT', {
+        providedIn: 'root',
+        factory: () => 'init',
+      });
+      const TOKEN_DESTROY = new InjectionToken('TOKEN_DESTROY', {
         providedIn: 'root',
-        factory: () => 'ngrx',
+        factory: () => 'destroy',
       });
       const Store = signalStore(
+        withState({ name: 'NgRx Store' }),
+        withHooks((store) => {
+          const tokenInit = inject(TOKEN_INIT);
+          const tokenDestroy = inject(TOKEN_DESTROY);
+          return {
+            onInit() {
+              messages.push(`${tokenInit} ${store.name()}`);
+            },
+            onDestroy() {
+              messages.push(`${tokenDestroy} ${store.name()}`);
+            },
+          };
+        })
+      );
+      const { destroy } = createLocalService(Store);
+
+      expect(messages).toEqual(['init NgRx Store']);
+
+      destroy();
+      expect(messages).toEqual(['init NgRx Store', 'destroy NgRx Store']);
+    });
+
+    it('executes hooks without injection context', () => {
+      const messages: string[] = [];
+      const Store = signalStore(
+        withState({ name: 'NgRx Store' }),
         withHooks({
-          onInit() {
-            inject(TOKEN);
-            messages.push('onInit');
+          onInit(store) {
+            messages.push(`init ${store.name()}`);
           },
-          onDestroy() {
-            // inject(TOKEN);
-            messages.push('onDestroy');
+          onDestroy(store) {
+            messages.push(`destroy ${store.name()}`);
           },
         })
       );
       const { destroy } = createLocalService(Store);
 
-      expect(messages).toEqual(['onInit']);
+      expect(messages).toEqual(['init NgRx Store']);
 
       destroy();
-      expect(messages).toEqual(['onInit', 'onDestroy']);
+      expect(messages).toEqual(['init NgRx Store', 'destroy NgRx Store']);
     });
 
     it('succeeds with onDestroy and providedIn: root', () => {

modules/signals/src/with-hooks.ts

@@ -7,7 +7,7 @@ import {
 } from './signal-store-models';
 import { Prettify } from './ts-helpers';
 
-type HooksFactory<Input extends SignalStoreFeatureResult> = (
+type HookFn<Input extends SignalStoreFeatureResult> = (
   store: Prettify<
     SignalStoreSlices<Input['state']> &
       Input['signals'] &
@@ -16,11 +16,45 @@ type HooksFactory<Input extends SignalStoreFeatureResult> = (
   >
 ) => void;
 
+type HooksFactory<Input extends SignalStoreFeatureResult> = (
+  store: Prettify<
+    SignalStoreSlices<Input['state']> &
+      Input['signals'] &
+      Input['methods'] &
+      StateSignal<Prettify<Input['state']>>
+  >
+) => {
+  onInit?: () => void;
+  onDestroy?: () => void;
+};
+
 export function withHooks<Input extends SignalStoreFeatureResult>(hooks: {
-  onInit?: HooksFactory<Input>;
-  onDestroy?: HooksFactory<Input>;
-}): SignalStoreFeature<Input, EmptyFeatureResult> {
+  onInit?: HookFn<Input>;
+  onDestroy?: HookFn<Input>;
+}): SignalStoreFeature<Input, EmptyFeatureResult>;
+export function withHooks<Input extends SignalStoreFeatureResult>(
+  hooks: HooksFactory<Input>
+): SignalStoreFeature<Input, EmptyFeatureResult>;
+
+export function withHooks<Input extends SignalStoreFeatureResult>(
+  hooksOrFactory:
+    | {
+        onInit?: HookFn<Input>;
+        onDestroy?: HookFn<Input>;
+      }
+    | HooksFactory<Input>
+): SignalStoreFeature<Input, EmptyFeatureResult> {
   return (store) => {
+    const storeProps = {
+      [STATE_SIGNAL]: store[STATE_SIGNAL],
+      ...store.slices,
+      ...store.signals,
+      ...store.methods,
+    };
+    const hooks =
+      typeof hooksOrFactory === 'function'
+        ? hooksOrFactory(storeProps)
+        : hooksOrFactory;
     const createHook = (name: keyof typeof hooks) => {
       const hook = hooks[name];
       const currentHook = store.hooks[name];
@@ -31,12 +65,7 @@ export function withHooks<Input extends SignalStoreFeatureResult>(hooks: {
               currentHook();
             }
 
-            hook({
-              [STATE_SIGNAL]: store[STATE_SIGNAL],
-              ...store.slices,
-              ...store.signals,
-              ...store.methods,
-            });
+            hook(storeProps);
           }
         : currentHook;
     };