fix: allow chaining of onCreate, closes #1

mfellner · Jul 31, 2022 · f9b1acc · f9b1acc
1 parent 272c14d
commit f9b1acc
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 15 deletions.
diff --git a/lib/store-factory.test.ts b/lib/store-factory.test.ts
@@ -183,7 +183,7 @@ describe('store-factory', () => {
     expect(context.fn).toHaveBeenCalledTimes(2);
   });
 
-  test('onCreate', () => {
+  test('onCreate passes the proxy state object', () => {
     let count = 0;
 
     const state = createFactory({ count: 0 })
@@ -208,6 +208,27 @@ describe('store-factory', () => {
     expect(count).toBe(1);
   });
 
+  test('onCreate can be chained and unsubscriptions will be called', () => {
+    const unsub1 = jest.fn();
+    const unsub2 = jest.fn();
+
+    const fn1 = jest.fn(() => unsub1);
+    const fn2 = jest.fn(() => unsub2);
+
+    const state = createFactory({ count: 0 })
+      .onCreate(() => fn1())
+      .onCreate(() => fn2())
+      .create();
+
+    expect(fn1).toHaveBeenCalledTimes(1);
+    expect(fn2).toHaveBeenCalledTimes(1);
+
+    state.$unsubscribe();
+
+    expect(unsub1).toHaveBeenCalledTimes(1);
+    expect(unsub2).toHaveBeenCalledTimes(1);
+  });
+
   test('$unsubscribe unsubscribes all subscriptions', () => {
     let count1 = 0;
     let count2 = 0;

diff --git a/lib/store-factory.ts b/lib/store-factory.ts
@@ -36,7 +36,7 @@ export function createFactory<State extends {}, Context = void>(
     baseDerivedProps: {},
     baseSubscriptions: [],
     unsubscriptions: [],
-    onCreate: () => undefined,
+    onCreateFns: [],
   });
 }
 
@@ -210,14 +210,14 @@ function factory<S extends {}, C extends {}, A extends Actions<S & U, C>, U exte
   baseDerivedProps,
   baseSubscriptions,
   unsubscriptions,
-  onCreate,
+  onCreateFns,
 }: {
   baseState: S;
   baseActions: A;
   baseDerivedProps: DerivedProps<S, U>;
   baseSubscriptions: Subscription<S, C>[];
   unsubscriptions: Array<() => void>;
-  onCreate: OnCreateFn<S, C, A, U>;
+  onCreateFns: OnCreateFn<S, C, A, U>[];
 }): Factory<S, C, A, U> {
   return {
     [isFactoryProp]: true,
@@ -229,7 +229,7 @@ function factory<S extends {}, C extends {}, A extends Actions<S & U, C>, U exte
         baseDerivedProps,
         baseSubscriptions,
         unsubscriptions,
-        onCreate,
+        onCreateFns,
       });
     },
 
@@ -240,7 +240,7 @@ function factory<S extends {}, C extends {}, A extends Actions<S & U, C>, U exte
         baseDerivedProps: combineDerivedProps(baseDerivedProps, derivedProps),
         baseSubscriptions,
         unsubscriptions,
-        onCreate,
+        onCreateFns,
       });
     },
 
@@ -251,7 +251,7 @@ function factory<S extends {}, C extends {}, A extends Actions<S & U, C>, U exte
         baseDerivedProps,
         baseSubscriptions: [...baseSubscriptions, [subscription, ...args]],
         unsubscriptions,
-        onCreate,
+        onCreateFns,
       });
     },
 
@@ -272,7 +272,7 @@ function factory<S extends {}, C extends {}, A extends Actions<S & U, C>, U exte
           ],
         ],
         unsubscriptions,
-        onCreate,
+        onCreateFns,
       });
     },
 
@@ -283,7 +283,7 @@ function factory<S extends {}, C extends {}, A extends Actions<S & U, C>, U exte
         baseDerivedProps,
         baseSubscriptions,
         unsubscriptions,
-        onCreate: fn,
+        onCreateFns: [...onCreateFns, fn],
       });
     },
 
@@ -341,13 +341,13 @@ function factory<S extends {}, C extends {}, A extends Actions<S & U, C>, U exte
         unsubscriptions.push(unsubscription);
       }
 
-      const unsubscribeFn = onCreate(derivedProxy as FactoryResult<S, C, A, U>);
+      const unsubscribeFns = onCreateFns
+        .map((fn) => fn(derivedProxy as FactoryResult<S, C, A, U>))
+        .filter((fn): fn is UnsubscribeFn => Boolean(fn));
 
-      if (unsubscribeFn) {
-        // Add the unsubscribe function returned by onCreate to the rest of the
-        // unsubscribe functions. They will be called when `state.$unsubscribe()` is called.
-        unsubscriptions.push(unsubscribeFn);
-      }
+      // Add the unsubscribe functions returned by onCreate functions to the rest of the
+      // unsubscribe functions. They will be called when `state.$unsubscribe()` is called.
+      unsubscriptions.push(...unsubscribeFns);
 
       return derivedProxy as FactoryResult<S, C, A, U>;
     },