fix(create-injection-token): allows multi token to work correctly

ngxtension · Oct 5, 2023 · 1f2915e · 1f2915e
1 parent 9ce8825
commit 1f2915e
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 35 deletions.
diff --git a/libs/ngxtension/create-injection-token/src/create-injection-token.spec.ts b/libs/ngxtension/create-injection-token/src/create-injection-token.spec.ts
@@ -101,4 +101,19 @@ describe(createInjectionToken.name, () => {
 			expect(value).toEqual(1);
 		});
 	});
+
+	describe('given multi injection token', () => {
+		const [injectFn, provideFn] = createInjectionToken(() => 1, {
+			multi: true,
+		});
+		it('then return value as array', () => {
+			TestBed.configureTestingModule({
+				providers: [provideFn(), provideFn(2)],
+			}).runInInjectionContext(() => {
+				const values = injectFn();
+				expect(Array.isArray(values)).toEqual(true);
+				expect(values).toEqual([1, 2]);
+			});
+		});
+	});
 });
diff --git a/libs/ngxtension/create-injection-token/src/create-injection-token.ts b/libs/ngxtension/create-injection-token/src/create-injection-token.ts
@@ -39,22 +39,23 @@ export type CreateInjectionTokenOptions<
 > =
 	// this means TFunction has no arguments
 	(TFactoryDeps[0] extends undefined
-		? {
-				isRoot: boolean;
-				deps?: never;
-		  }
-		: {
-				isRoot?: boolean;
-				deps: CreateInjectionTokenDeps<TFactory, TFactoryDeps>;
-		  }) & {
+		? { deps?: never }
+		: { deps: CreateInjectionTokenDeps<TFactory, TFactoryDeps> }) & {
+		isRoot?: boolean;
+		multi?: boolean;
 		token?: InjectionToken<ReturnType<TFactory>>;
 		extraProviders?: Provider | EnvironmentProviders;
 	};
 
-type InjectFn<
+type CreateProvideFnOptions<
 	TFactory extends (...args: any[]) => any,
-	TFactoryReturn extends ReturnType<TFactory> = ReturnType<TFactory>
-> = {
+	TFactoryDeps extends Parameters<TFactory> = Parameters<TFactory>
+> = Pick<
+	CreateInjectionTokenOptions<TFactory, TFactoryDeps>,
+	'deps' | 'extraProviders' | 'multi'
+>;
+
+type InjectFn<TFactoryReturn> = {
 	(): TFactoryReturn;
 	(
 		injectOptions: InjectOptions & { optional?: false } & {
@@ -66,12 +67,13 @@ type InjectFn<
 	): TFactoryReturn | null;
 };
 
-export type CreateInjectionTokenReturn<
-	TFactory extends (...args: any[]) => any,
-	TFactoryReturn extends ReturnType<TFactory> = ReturnType<TFactory>
-> = [
-	InjectFn<TFactory, TFactoryReturn>,
-	(value?: TFactoryReturn) => Provider,
+export type CreateInjectionTokenReturn<TFactoryReturn> = [
+	InjectFn<TFactoryReturn>,
+	(
+		value?: TFactoryReturn extends Array<infer TReturn>
+			? TReturn
+			: TFactoryReturn
+	) => Provider,
 	InjectionToken<TFactoryReturn>
 ];
 
@@ -94,22 +96,23 @@ function createProvideFn<
 >(
 	token: InjectionToken<TValue>,
 	factory: (...args: any[]) => TValue,
-	deps?: CreateInjectionTokenDeps<TFactory, TFactoryDeps>,
-	extraProviders?: Provider | EnvironmentProviders
+	opts: CreateProvideFnOptions<TFactory, TFactoryDeps> = {}
 ) {
+	const { deps = [], multi = false, extraProviders = [] } = opts;
 	return (value?: TValue) => {
 		let provider: Provider;
 		if (value) {
-			provider = { provide: token, useValue: value };
+			provider = { provide: token, useValue: value, multi };
 		} else {
 			provider = {
 				provide: token,
 				useFactory: factory,
-				deps: (deps ?? []) as FactoryProvider['deps'],
+				deps: deps as FactoryProvider['deps'],
+				multi,
 			};
 		}
 
-		return extraProviders ? [extraProviders, provider] : provider;
+		return [extraProviders, provider];
 	};
 }
 
@@ -133,11 +136,17 @@ function createProvideFn<
 export function createInjectionToken<
 	TFactory extends (...args: any[]) => any,
 	TFactoryDeps extends Parameters<TFactory> = Parameters<TFactory>,
-	TFactoryReturn extends ReturnType<TFactory> = ReturnType<TFactory>
+	TOptions extends CreateInjectionTokenOptions<
+		TFactory,
+		TFactoryDeps
+	> = CreateInjectionTokenOptions<TFactory, TFactoryDeps>,
+	TFactoryReturn = TOptions['multi'] extends true
+		? Array<ReturnType<TFactory>>
+		: ReturnType<TFactory>
 >(
 	factory: TFactory,
-	options?: CreateInjectionTokenOptions<TFactory, TFactoryDeps>
-): CreateInjectionTokenReturn<TFactory, TFactoryReturn> {
+	options?: TOptions
+): CreateInjectionTokenReturn<TFactoryReturn> {
 	const tokenName = factory.name || factory.toString();
 	const opts =
 		options ??
@@ -162,23 +171,25 @@ createInjectionToken is creating a root InjectionToken but an external token is
 		});
 
 		return [
-			createInjectFn(token) as CreateInjectionTokenReturn<
-				TFactory,
-				TFactoryReturn
-			>[0],
-			createProvideFn(token, factory, opts.deps),
+			createInjectFn(token) as CreateInjectionTokenReturn<TFactoryReturn>[0],
+			createProvideFn(
+				token,
+				factory,
+				opts as CreateProvideFnOptions<TFactory, TFactoryDeps>
+			),
 			token,
 		];
 	}
 
 	const token =
 		opts.token || new InjectionToken<TFactoryReturn>(`Token for ${tokenName}`);
 	return [
-		createInjectFn(token) as CreateInjectionTokenReturn<
-			TFactory,
-			TFactoryReturn
-		>[0],
-		createProvideFn(token, factory, opts.deps, opts.extraProviders),
+		createInjectFn(token) as CreateInjectionTokenReturn<TFactoryReturn>[0],
+		createProvideFn(
+			token,
+			factory,
+			opts as CreateProvideFnOptions<TFactory, TFactoryDeps>
+		),
 		token,
 	];
 }