diInit receive optional context

darky · May 18, 2023 · 4d38cc0 · 4d38cc0
1 parent ba2ac9f
commit 4d38cc0
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 4 deletions.
diff --git a/index.ts b/index.ts
@@ -1,11 +1,13 @@
 import { AsyncLocalStorage } from 'async_hooks';
 
-export const als = new AsyncLocalStorage<{
+type AlsContext = {
   deps: Map<unknown, unknown>;
   once: Map<unknown, unknown>;
   state: Map<unknown, unknown>;
   [k: string]: unknown;
-}>();
+}
+
+export const als = new AsyncLocalStorage<AlsContext>();
 
 const globalState = new Map<unknown, unknown>();
 
@@ -61,8 +63,8 @@ export const diHas = <T>(dep: T | string): boolean => {
   return (store.deps as Map<T | string, T>).has(dep);
 };
 
-export const diInit = <T>(cb: () => T) => {
-  return diExists() ? cb() : als.run({ deps: new Map(), once: new Map(), state: new Map() }, cb);
+export const diInit = <T>(cb: () => T, ctx?: AlsContext) => {
+  return diExists() ? cb() : als.run(ctx ?? { deps: new Map(), once: new Map(), state: new Map() }, cb);
 };
 
 export const diOnce = <T extends Function>(fn: T): T => {

diff --git a/test.ts b/test.ts
@@ -191,6 +191,17 @@ test('diInit on existing als', async () => {
   })
 });
 
+test('diInit can receive context', async () => {
+  const ctx = {
+    deps: new Map([['foo', 'bar']]),
+    state: new Map(),
+    once: new Map()
+  }
+  await diInit(async () => {
+    assert.equal(diDep('foo'), 'bar')
+  }, ctx)
+});
+
 test('diExists - false', () => {
   assert.equal(diExists(), false);
 });