From 3c80b580dcb565bd9ee0a7a03ee5c6a7a6776a32 Mon Sep 17 00:00:00 2001
From: Arnoud de Vries <6420061+arnoud-dv@users.noreply.github.com>
Date: Thu, 8 May 2025 21:13:41 +0200
Subject: [PATCH] feature(angular-query): support providing QueryClient via an
 InjectionToken

---
 .../__tests__/provide-query-client.test.ts    | 41 +++++++++++++++++++
 .../__tests__/provide-tanstack-query.test.ts  | 41 +++++++++++++++++++
 .../src/providers.ts                          | 35 +++++++++++-----
 3 files changed, 107 insertions(+), 10 deletions(-)
 create mode 100644 packages/angular-query-experimental/src/__tests__/provide-query-client.test.ts
 create mode 100644 packages/angular-query-experimental/src/__tests__/provide-tanstack-query.test.ts

diff --git a/packages/angular-query-experimental/src/__tests__/provide-query-client.test.ts b/packages/angular-query-experimental/src/__tests__/provide-query-client.test.ts
new file mode 100644
index 00000000000..5758a17457f
--- /dev/null
+++ b/packages/angular-query-experimental/src/__tests__/provide-query-client.test.ts
@@ -0,0 +1,41 @@
+import { TestBed } from '@angular/core/testing'
+import { describe, expect, test } from 'vitest'
+import {
+  InjectionToken,
+  provideExperimentalZonelessChangeDetection,
+} from '@angular/core'
+import { QueryClient } from '@tanstack/query-core'
+import { provideQueryClient } from '../providers'
+
+describe('provideQueryClient', () => {
+  test('should provide a QueryClient instance directly', () => {
+    const queryClient = new QueryClient()
+
+    TestBed.configureTestingModule({
+      providers: [
+        provideExperimentalZonelessChangeDetection(),
+        provideQueryClient(queryClient),
+      ],
+    })
+
+    const providedQueryClient = TestBed.inject(QueryClient)
+    expect(providedQueryClient).toBe(queryClient)
+  })
+
+  test('should provide a QueryClient instance using an InjectionToken', () => {
+    const queryClient = new QueryClient()
+    const CUSTOM_QUERY_CLIENT = new InjectionToken<QueryClient>('', {
+      factory: () => queryClient,
+    })
+
+    TestBed.configureTestingModule({
+      providers: [
+        provideExperimentalZonelessChangeDetection(),
+        provideQueryClient(CUSTOM_QUERY_CLIENT),
+      ],
+    })
+
+    const providedQueryClient = TestBed.inject(QueryClient)
+    expect(providedQueryClient).toBe(queryClient)
+  })
+})
diff --git a/packages/angular-query-experimental/src/__tests__/provide-tanstack-query.test.ts b/packages/angular-query-experimental/src/__tests__/provide-tanstack-query.test.ts
new file mode 100644
index 00000000000..c35eb406a5c
--- /dev/null
+++ b/packages/angular-query-experimental/src/__tests__/provide-tanstack-query.test.ts
@@ -0,0 +1,41 @@
+import { TestBed } from '@angular/core/testing'
+import { describe, expect, test } from 'vitest'
+import {
+  InjectionToken,
+  provideExperimentalZonelessChangeDetection,
+} from '@angular/core'
+import { QueryClient } from '@tanstack/query-core'
+import { provideTanStackQuery } from '../providers'
+
+describe('provideTanStackQuery', () => {
+  test('should provide a QueryClient instance directly', () => {
+    const queryClient = new QueryClient()
+
+    TestBed.configureTestingModule({
+      providers: [
+        provideExperimentalZonelessChangeDetection(),
+        provideTanStackQuery(queryClient),
+      ],
+    })
+
+    const providedQueryClient = TestBed.inject(QueryClient)
+    expect(providedQueryClient).toBe(queryClient)
+  })
+
+  test('should provide a QueryClient instance using an InjectionToken', () => {
+    const queryClient = new QueryClient()
+    const CUSTOM_QUERY_CLIENT = new InjectionToken<QueryClient>('', {
+      factory: () => queryClient,
+    })
+
+    TestBed.configureTestingModule({
+      providers: [
+        provideExperimentalZonelessChangeDetection(),
+        provideTanStackQuery(CUSTOM_QUERY_CLIENT),
+      ],
+    })
+
+    const providedQueryClient = TestBed.inject(QueryClient)
+    expect(providedQueryClient).toBe(queryClient)
+  })
+})
diff --git a/packages/angular-query-experimental/src/providers.ts b/packages/angular-query-experimental/src/providers.ts
index 98ea2e04ed8..0a7f784f4c0 100644
--- a/packages/angular-query-experimental/src/providers.ts
+++ b/packages/angular-query-experimental/src/providers.ts
@@ -1,6 +1,7 @@
 import {
   DestroyRef,
   ENVIRONMENT_INITIALIZER,
+  InjectionToken,
   PLATFORM_ID,
   computed,
   effect,
@@ -25,10 +26,19 @@ import type {
  * for the entire application. You can use `provideQueryClient` to provide a
  * different `QueryClient` instance for a part of the application.
  * @param queryClient - the `QueryClient` instance to provide.
- * @public
+ * @returns a provider object that can be used to provide the `QueryClient` instance.
  */
-export function provideQueryClient(queryClient: QueryClient) {
-  return { provide: QueryClient, useValue: queryClient }
+export function provideQueryClient(
+  queryClient: QueryClient | InjectionToken<QueryClient>,
+): Provider {
+  return {
+    provide: QueryClient,
+    useFactory: () => {
+      return queryClient instanceof InjectionToken
+        ? inject(queryClient)
+        : queryClient
+    },
+  }
 }
 
 /**
@@ -83,15 +93,14 @@ export function provideQueryClient(queryClient: QueryClient) {
  *   }
  * )
  * ```
- * @param queryClient - A `QueryClient` instance.
+ * @param queryClient - A `QueryClient` instance, or an `InjectionToken` which provides a `QueryClient`.
  * @param features - Optional features to configure additional Query functionality.
  * @returns A set of providers to set up TanStack Query.
- * @public
  * @see https://tanstack.com/query/v5/docs/framework/angular/quick-start
  * @see withDevtools
  */
 export function provideTanStackQuery(
-  queryClient: QueryClient,
+  queryClient: QueryClient | InjectionToken<QueryClient>,
   ...features: Array<QueryFeatures>
 ): EnvironmentProviders {
   return makeEnvironmentProviders([
@@ -100,10 +109,16 @@ export function provideTanStackQuery(
       // Do not use provideEnvironmentInitializer while Angular < v19 is supported
       provide: ENVIRONMENT_INITIALIZER,
       multi: true,
-      useValue: () => {
-        queryClient.mount()
-        // Unmount the query client on application destroy
-        inject(DestroyRef).onDestroy(() => queryClient.unmount())
+      useFactory: () => {
+        const client =
+          queryClient instanceof InjectionToken
+            ? inject(queryClient)
+            : queryClient
+        return () => {
+          client.mount()
+          // Unmount the query client on application destroy
+          inject(DestroyRef).onDestroy(() => client.unmount())
+        }
       },
     },
     features.map((feature) => feature.ɵproviders),