From 47ea64f17244ce1dc605bcb77982cdb6e8a55c72 Mon Sep 17 00:00:00 2001
From: Brandon Roberts <robertsbt@gmail.com>
Date: Tue, 26 Apr 2022 21:23:26 -0500
Subject: [PATCH] chore: add tests

---
 .../spec/component-store.spec.ts              | 86 ++++++++++++++++++-
 .../component-store/src/component-store.ts    |  4 +-
 2 files changed, 87 insertions(+), 3 deletions(-)

diff --git a/modules/component-store/spec/component-store.spec.ts b/modules/component-store/spec/component-store.spec.ts
index 2aaa3f54e9..34a00bb8c5 100644
--- a/modules/component-store/spec/component-store.spec.ts
+++ b/modules/component-store/spec/component-store.spec.ts
@@ -1,5 +1,9 @@
 /* eslint-disable @typescript-eslint/no-non-null-assertion */
-import { ComponentStore } from '@ngrx/component-store';
+import {
+  ComponentStore,
+  OnStateInit,
+  OnStoreInit,
+} from '@ngrx/component-store';
 import { fakeSchedulers, marbles } from 'rxjs-marbles/jest';
 import {
   of,
@@ -1447,4 +1451,84 @@ 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';
+    let logs: string[] = [];
+    class LifecycleStore
+      extends ComponentStore<LifeCycle>
+      implements OnStoreInit, OnStateInit
+    {
+      constructor(state?: LifeCycle) {
+        super(state);
+      }
+
+      logEffect = this.effect(
+        tap<void>(() => {
+          logs.push('effect');
+        })
+      );
+
+      ngrxOnStoreInit() {
+        logs.push(onStoreInitMessage);
+      }
+
+      ngrxOnStateInit() {
+        logs.push(onStateInitMessage);
+      }
+    }
+
+    let componentStore: LifecycleStore;
+
+    beforeEach(() => {
+      logs = [];
+    });
+
+    it('should call the OnInitStore lifecycle hook if defined', async () => {
+      componentStore = new LifecycleStore({ init: true });
+
+      expect(logs[0]).toBe(onStoreInitMessage);
+    });
+
+    it('should only call the OnInitStore lifecycle hook once', async () => {
+      componentStore = new LifecycleStore({ init: true });
+      expect(logs[0]).toBe(onStoreInitMessage);
+
+      logs = [];
+      componentStore.setState({ init: false });
+
+      expect(logs.length).toBe(0);
+    });
+
+    it('should call the OnInitState lifecycle hook if defined and state is set eagerly', async () => {
+      componentStore = new LifecycleStore({ init: true });
+
+      expect(logs[1]).toBe(onStateInitMessage);
+    });
+
+    it('should call the OnInitState lifecycle hook if defined and after state is set lazily', async () => {
+      componentStore = new LifecycleStore();
+
+      expect(logs.length).toBe(1);
+
+      componentStore.setState({ init: true });
+
+      expect(logs[1]).toBe(onStateInitMessage);
+    });
+
+    it('should only call the OnInitStore lifecycle hook once', async () => {
+      componentStore = new LifecycleStore({ init: true });
+
+      expect(logs[1]).toBe(onStateInitMessage);
+      logs = [];
+      componentStore.setState({ init: false });
+
+      expect(logs.length).toBe(0);
+    });
+  });
 });
diff --git a/modules/component-store/src/component-store.ts b/modules/component-store/src/component-store.ts
index 53b7ecaac8..d0141b8be2 100644
--- a/modules/component-store/src/component-store.ts
+++ b/modules/component-store/src/component-store.ts
@@ -71,7 +71,7 @@ export class ComponentStore<T extends object> implements OnDestroy {
 
   constructor(@Optional() @Inject(INITIAL_STATE_TOKEN) defaultState?: T) {
     // check/call store init hook
-    this.callInitStoreHook();
+    this.callInitStoreHook(defaultState);
 
     // State can be initialized either through constructor or setState.
     if (defaultState) {
@@ -79,7 +79,7 @@ export class ComponentStore<T extends object> implements OnDestroy {
     }
   }
 
-  private callInitStoreHook() {
+  private callInitStoreHook(ds?: T) {
     const onStoreInit: Function | undefined = (
       this as unknown as ComponentStore<T> & OnStoreInit
     )['ngrxOnStoreInit'];