From 2d5a3df4e6bcf735f804e7fad868aedb47289977 Mon Sep 17 00:00:00 2001
From: Shahar Har-Shuv <sharshuv@quotient.com>
Date: Wed, 1 May 2024 17:30:09 -0400
Subject: [PATCH] feat(angular-query-experimental): run query options callbacks
 in injection context

This is to allow consumers to use `inject` inside callbacks like `queryFn`, `initialData` etc.
---
 .../src/__tests__/inject-query.test.ts        | 419 ++++++++++--------
 .../src/inject-query.ts                       |  37 +-
 .../injection-context-proxy.ts                |  37 ++
 3 files changed, 287 insertions(+), 206 deletions(-)
 create mode 100644 packages/angular-query-experimental/src/util/injection-context-proxy/injection-context-proxy.ts

diff --git a/packages/angular-query-experimental/src/__tests__/inject-query.test.ts b/packages/angular-query-experimental/src/__tests__/inject-query.test.ts
index 5ba1ac4085..a61af9cf2e 100644
--- a/packages/angular-query-experimental/src/__tests__/inject-query.test.ts
+++ b/packages/angular-query-experimental/src/__tests__/inject-query.test.ts
@@ -5,345 +5,386 @@ import {
   inject,
   input,
   signal,
-} from '@angular/core'
-import { TestBed, fakeAsync, flush, tick } from '@angular/core/testing'
-import { QueryClient } from '@tanstack/query-core'
-import { describe, expect, vi } from 'vitest'
-import { injectQuery } from '../inject-query'
-import { provideAngularQuery } from '../providers'
+} from "@angular/core";
+import { TestBed, fakeAsync, flush, tick } from "@angular/core/testing";
+import { QueryClient } from "@tanstack/query-core";
+import { describe, expect, vi } from "vitest";
+import { injectQuery } from "../inject-query";
+import { provideAngularQuery } from "../providers";
 import {
   delayedFetcher,
   getSimpleFetcherWithReturnData,
   rejectFetcher,
   setSignalInputs,
   simpleFetcher,
-} from './test-utils'
+} from "./test-utils";
 
-describe('injectQuery', () => {
+describe("injectQuery", () => {
   beforeEach(() => {
     TestBed.configureTestingModule({
       providers: [provideAngularQuery(new QueryClient())],
-    })
-  })
+    });
+  });
 
-  test('should return pending status initially', fakeAsync(() => {
+  test("should return pending status initially", fakeAsync(() => {
     const query = TestBed.runInInjectionContext(() => {
       return injectQuery(() => ({
-        queryKey: ['key1'],
+        queryKey: ["key1"],
         queryFn: simpleFetcher,
-      }))
-    })
+      }));
+    });
 
-    expect(query.status()).toBe('pending')
-    expect(query.isPending()).toBe(true)
-    expect(query.isFetching()).toBe(true)
-    expect(query.isStale()).toBe(true)
-    expect(query.isFetched()).toBe(false)
+    expect(query.status()).toBe("pending");
+    expect(query.isPending()).toBe(true);
+    expect(query.isFetching()).toBe(true);
+    expect(query.isStale()).toBe(true);
+    expect(query.isFetched()).toBe(false);
 
-    flush()
-  }))
+    flush();
+  }));
 
-  test('should resolve to success and update signal: injectQuery()', fakeAsync(() => {
+  test("should resolve to success and update signal: injectQuery()", fakeAsync(() => {
     const query = TestBed.runInInjectionContext(() => {
       return injectQuery(() => ({
-        queryKey: ['key2'],
-        queryFn: getSimpleFetcherWithReturnData('result2'),
-      }))
-    })
-
-    flush()
-
-    expect(query.status()).toBe('success')
-    expect(query.data()).toBe('result2')
-    expect(query.isPending()).toBe(false)
-    expect(query.isFetching()).toBe(false)
-    expect(query.isFetched()).toBe(true)
-    expect(query.isSuccess()).toBe(true)
-  }))
-
-  test('should reject and update signal', fakeAsync(() => {
+        queryKey: ["key2"],
+        queryFn: getSimpleFetcherWithReturnData("result2"),
+      }));
+    });
+
+    flush();
+
+    expect(query.status()).toBe("success");
+    expect(query.data()).toBe("result2");
+    expect(query.isPending()).toBe(false);
+    expect(query.isFetching()).toBe(false);
+    expect(query.isFetched()).toBe(true);
+    expect(query.isSuccess()).toBe(true);
+  }));
+
+  test("should reject and update signal", fakeAsync(() => {
     const query = TestBed.runInInjectionContext(() => {
       return injectQuery(() => ({
         retry: false,
-        queryKey: ['key3'],
+        queryKey: ["key3"],
         queryFn: rejectFetcher,
-      }))
-    })
+      }));
+    });
 
-    flush()
+    flush();
 
-    expect(query.status()).toBe('error')
-    expect(query.data()).toBe(undefined)
-    expect(query.error()).toMatchObject({ message: 'Some error' })
-    expect(query.isPending()).toBe(false)
-    expect(query.isFetching()).toBe(false)
-    expect(query.isError()).toBe(true)
-    expect(query.failureCount()).toBe(1)
-    expect(query.failureReason()).toMatchObject({ message: 'Some error' })
-  }))
+    expect(query.status()).toBe("error");
+    expect(query.data()).toBe(undefined);
+    expect(query.error()).toMatchObject({ message: "Some error" });
+    expect(query.isPending()).toBe(false);
+    expect(query.isFetching()).toBe(false);
+    expect(query.isError()).toBe(true);
+    expect(query.failureCount()).toBe(1);
+    expect(query.failureReason()).toMatchObject({ message: "Some error" });
+  }));
 
-  test('should update query on options contained signal change', fakeAsync(() => {
-    const key = signal(['key6', 'key7'])
-    const spy = vi.fn(simpleFetcher)
+  test("should update query on options contained signal change", fakeAsync(() => {
+    const key = signal(["key6", "key7"]);
+    const spy = vi.fn(simpleFetcher);
 
     const query = TestBed.runInInjectionContext(() => {
       return injectQuery(() => ({
         queryKey: key(),
         queryFn: spy,
-      }))
-    })
-    flush()
-    expect(spy).toHaveBeenCalledTimes(1)
+      }));
+    });
+    flush();
+    expect(spy).toHaveBeenCalledTimes(1);
 
-    expect(query.status()).toBe('success')
+    expect(query.status()).toBe("success");
 
-    key.set(['key8'])
-    TestBed.flushEffects()
+    key.set(["key8"]);
+    TestBed.flushEffects();
 
-    expect(spy).toHaveBeenCalledTimes(2)
+    expect(spy).toHaveBeenCalledTimes(2);
 
-    flush()
-  }))
+    flush();
+  }));
 
-  test('should only run query once enabled signal is set to true', fakeAsync(() => {
-    const spy = vi.fn(simpleFetcher)
-    const enabled = signal(false)
+  test("should only run query once enabled signal is set to true", fakeAsync(() => {
+    const spy = vi.fn(simpleFetcher);
+    const enabled = signal(false);
 
     const query = TestBed.runInInjectionContext(() => {
       return injectQuery(() => ({
-        queryKey: ['key9'],
+        queryKey: ["key9"],
         queryFn: spy,
         enabled: enabled(),
-      }))
-    })
+      }));
+    });
 
-    expect(spy).not.toHaveBeenCalled()
-    expect(query.status()).toBe('pending')
+    expect(spy).not.toHaveBeenCalled();
+    expect(query.status()).toBe("pending");
 
-    enabled.set(true)
-    TestBed.flushEffects()
-    flush()
-    expect(spy).toHaveBeenCalledTimes(1)
-    expect(query.status()).toBe('success')
-  }))
+    enabled.set(true);
+    TestBed.flushEffects();
+    flush();
+    expect(spy).toHaveBeenCalledTimes(1);
+    expect(query.status()).toBe("success");
+  }));
 
-  test('should properly execute dependant queries', fakeAsync(() => {
+  test("should properly execute dependant queries", fakeAsync(() => {
     const query1 = TestBed.runInInjectionContext(() => {
       return injectQuery(() => ({
-        queryKey: ['dependant1'],
+        queryKey: ["dependant1"],
         queryFn: simpleFetcher,
-      }))
-    })
+      }));
+    });
 
-    const dependentQueryFn = vi.fn().mockImplementation(delayedFetcher(1000))
+    const dependentQueryFn = vi.fn().mockImplementation(delayedFetcher(1000));
 
     const query2 = TestBed.runInInjectionContext(() => {
       return injectQuery(
         computed(() => ({
-          queryKey: ['dependant2'],
+          queryKey: ["dependant2"],
           queryFn: dependentQueryFn,
           enabled: !!query1.data(),
         })),
-      )
-    })
+      );
+    });
 
-    expect(query1.data()).toStrictEqual(undefined)
-    expect(query2.fetchStatus()).toStrictEqual('idle')
-    expect(dependentQueryFn).not.toHaveBeenCalled()
+    expect(query1.data()).toStrictEqual(undefined);
+    expect(query2.fetchStatus()).toStrictEqual("idle");
+    expect(dependentQueryFn).not.toHaveBeenCalled();
 
-    tick()
-    TestBed.flushEffects()
+    tick();
+    TestBed.flushEffects();
 
-    expect(query1.data()).toStrictEqual('Some data')
-    expect(query2.fetchStatus()).toStrictEqual('fetching')
+    expect(query1.data()).toStrictEqual("Some data");
+    expect(query2.fetchStatus()).toStrictEqual("fetching");
 
-    flush()
+    flush();
 
-    expect(query2.fetchStatus()).toStrictEqual('idle')
-    expect(query2.status()).toStrictEqual('success')
-    expect(dependentQueryFn).toHaveBeenCalledTimes(1)
+    expect(query2.fetchStatus()).toStrictEqual("idle");
+    expect(query2.status()).toStrictEqual("success");
+    expect(dependentQueryFn).toHaveBeenCalledTimes(1);
     expect(dependentQueryFn).toHaveBeenCalledWith(
-      expect.objectContaining({ queryKey: ['dependant2'] }),
-    )
-  }))
+      expect.objectContaining({ queryKey: ["dependant2"] }),
+    );
+  }));
 
-  test('should use the current value for the queryKey when refetch is called', fakeAsync(() => {
-    const fetchFn = vi.fn(simpleFetcher)
-    const keySignal = signal('key11')
+  test("should use the current value for the queryKey when refetch is called", fakeAsync(() => {
+    const fetchFn = vi.fn(simpleFetcher);
+    const keySignal = signal("key11");
 
     const query = TestBed.runInInjectionContext(() => {
       return injectQuery(() => ({
-        queryKey: ['key10', keySignal()],
+        queryKey: ["key10", keySignal()],
         queryFn: fetchFn,
         enabled: false,
-      }))
-    })
+      }));
+    });
 
-    expect(fetchFn).not.toHaveBeenCalled()
+    expect(fetchFn).not.toHaveBeenCalled();
 
     query.refetch().then(() => {
-      expect(fetchFn).toHaveBeenCalledTimes(1)
+      expect(fetchFn).toHaveBeenCalledTimes(1);
       expect(fetchFn).toHaveBeenCalledWith(
         expect.objectContaining({
-          queryKey: ['key10', 'key11'],
+          queryKey: ["key10", "key11"],
         }),
-      )
-    })
+      );
+    });
 
-    flush()
+    flush();
 
-    keySignal.set('key12')
+    keySignal.set("key12");
 
-    TestBed.flushEffects()
+    TestBed.flushEffects();
 
     query.refetch().then(() => {
-      expect(fetchFn).toHaveBeenCalledTimes(2)
+      expect(fetchFn).toHaveBeenCalledTimes(2);
       expect(fetchFn).toHaveBeenCalledWith(
         expect.objectContaining({
-          queryKey: ['key10', 'key12'],
+          queryKey: ["key10", "key12"],
         }),
-      )
-    })
+      );
+    });
 
-    flush()
-  }))
+    flush();
+  }));
 
-  describe('throwOnError', () => {
-    test('should evaluate throwOnError when query is expected to throw', fakeAsync(() => {
-      const boundaryFn = vi.fn()
+  describe("throwOnError", () => {
+    test("should evaluate throwOnError when query is expected to throw", fakeAsync(() => {
+      const boundaryFn = vi.fn();
       TestBed.runInInjectionContext(() => {
         return injectQuery(() => ({
-          queryKey: ['key12'],
+          queryKey: ["key12"],
           queryFn: rejectFetcher,
           throwOnError: boundaryFn,
-        }))
-      })
+        }));
+      });
 
-      flush()
+      flush();
 
-      expect(boundaryFn).toHaveBeenCalledTimes(1)
+      expect(boundaryFn).toHaveBeenCalledTimes(1);
       expect(boundaryFn).toHaveBeenCalledWith(
-        Error('Some error'),
+        Error("Some error"),
         expect.objectContaining({
-          state: expect.objectContaining({ status: 'error' }),
+          state: expect.objectContaining({ status: "error" }),
         }),
-      )
-    }))
+      );
+    }));
 
-    test('should throw when throwOnError is true', fakeAsync(() => {
+    test("should throw when throwOnError is true", fakeAsync(() => {
       TestBed.runInInjectionContext(() => {
         return injectQuery(() => ({
-          queryKey: ['key13'],
+          queryKey: ["key13"],
           queryFn: rejectFetcher,
           throwOnError: true,
-        }))
-      })
+        }));
+      });
 
       expect(() => {
-        flush()
-      }).toThrowError('Some error')
-      flush()
-    }))
+        flush();
+      }).toThrowError("Some error");
+      flush();
+    }));
 
-    test('should throw when throwOnError function returns true', fakeAsync(() => {
+    test("should throw when throwOnError function returns true", fakeAsync(() => {
       TestBed.runInInjectionContext(() => {
         return injectQuery(() => ({
-          queryKey: ['key14'],
+          queryKey: ["key14"],
           queryFn: rejectFetcher,
           throwOnError: () => true,
-        }))
-      })
+        }));
+      });
 
       expect(() => {
-        flush()
-      }).toThrowError('Some error')
-      flush()
-    }))
-  })
+        flush();
+      }).toThrowError("Some error");
+      flush();
+    }));
+  });
 
-  test('should set state to error when queryFn returns reject promise', fakeAsync(() => {
+  test("should set state to error when queryFn returns reject promise", fakeAsync(() => {
     const query = TestBed.runInInjectionContext(() => {
       return injectQuery(() => ({
         retry: false,
-        queryKey: ['key15'],
+        queryKey: ["key15"],
         queryFn: rejectFetcher,
-      }))
-    })
+      }));
+    });
 
-    expect(query.status()).toBe('pending')
+    expect(query.status()).toBe("pending");
 
-    flush()
+    flush();
 
-    expect(query.status()).toBe('error')
-  }))
+    expect(query.status()).toBe("error");
+  }));
 
-  test('should render with required signal inputs', fakeAsync(async () => {
+  test("should render with required signal inputs", fakeAsync(async () => {
     @Component({
-      selector: 'app-fake',
+      selector: "app-fake",
       template: `{{ query.data() }}`,
       standalone: true,
     })
     class FakeComponent {
-      name = input.required<string>()
+      name = input.required<string>();
 
       query = injectQuery(() => ({
-        queryKey: ['fake', this.name()],
+        queryKey: ["fake", this.name()],
         queryFn: () => Promise.resolve(this.name()),
-      }))
+      }));
     }
 
-    const fixture = TestBed.createComponent(FakeComponent)
+    const fixture = TestBed.createComponent(FakeComponent);
     setSignalInputs(fixture.componentInstance, {
-      name: 'signal-input-required-test',
-    })
+      name: "signal-input-required-test",
+    });
 
-    flush()
-    fixture.detectChanges()
+    flush();
+    fixture.detectChanges();
 
     expect(fixture.debugElement.nativeElement.textContent).toEqual(
-      'signal-input-required-test',
-    )
-  }))
+      "signal-input-required-test",
+    );
+  }));
 
-  test('should run options in injection context', fakeAsync(async () => {
+  test("should run options in injection context", fakeAsync(async () => {
     @Injectable()
     class FakeService {
       getData(name: string) {
-        return Promise.resolve(name)
+        return Promise.resolve(name);
       }
     }
 
     @Component({
-      selector: 'app-fake',
+      selector: "app-fake",
       template: `{{ query.data() }}`,
       standalone: true,
       providers: [FakeService],
     })
     class FakeComponent {
-      name = signal<string>('test name')
+      name = signal<string>("test name");
 
       query = injectQuery(() => {
-        const service = inject(FakeService)
+        const service = inject(FakeService);
 
         return {
-          queryKey: ['fake', this.name()],
+          queryKey: ["fake", this.name()],
+          queryFn: () => {
+            return service.getData(this.name());
+          },
+        };
+      });
+    }
+
+    const fixture = TestBed.createComponent(FakeComponent);
+    flush();
+    fixture.detectChanges();
+
+    expect(fixture.componentInstance.query.data()).toEqual("test name");
+
+    fixture.componentInstance.name.set("test name 2");
+    fixture.detectChanges();
+    flush();
+
+    expect(fixture.componentInstance.query.data()).toEqual("test name 2");
+  }));
+
+  test("should run queryFn in injection context", fakeAsync(async () => {
+    @Injectable()
+    class FakeService {
+      getData(name: string) {
+        return Promise.resolve(name);
+      }
+    }
+
+    @Component({
+      selector: "app-fake",
+      template: `{{ query.data() }}`,
+      standalone: true,
+      providers: [FakeService],
+    })
+    class FakeComponent {
+      name = signal<string>("test name");
+
+      query = injectQuery(() => {
+        return {
+          queryKey: ["fake", this.name()],
           queryFn: () => {
-            return service.getData(this.name())
+            const service = inject(FakeService);
+            return service.getData(this.name());
           },
-        }
-      })
+        };
+      });
     }
 
-    const fixture = TestBed.createComponent(FakeComponent)
-    flush()
-    fixture.detectChanges()
+    const fixture = TestBed.createComponent(FakeComponent);
+    flush();
+    fixture.detectChanges();
 
-    expect(fixture.componentInstance.query.data()).toEqual('test name')
+    expect(fixture.componentInstance.query.data()).toEqual("test name");
 
-    fixture.componentInstance.name.set('test name 2')
-    fixture.detectChanges()
-    flush()
+    fixture.componentInstance.name.set("test name 2");
+    fixture.detectChanges();
+    flush();
 
-    expect(fixture.componentInstance.query.data()).toEqual('test name 2')
-  }))
-})
+    expect(fixture.componentInstance.query.data()).toEqual("test name 2");
+  }));
+});
diff --git a/packages/angular-query-experimental/src/inject-query.ts b/packages/angular-query-experimental/src/inject-query.ts
index 22f6f5de2e..31ec8eef36 100644
--- a/packages/angular-query-experimental/src/inject-query.ts
+++ b/packages/angular-query-experimental/src/inject-query.ts
@@ -1,19 +1,20 @@
-import { QueryObserver } from '@tanstack/query-core'
-import { runInInjectionContext } from '@angular/core'
-import { assertInjector } from './util/assert-injector/assert-injector'
-import { injectQueryClient } from './inject-query-client'
-import { createBaseQuery } from './create-base-query'
-import type { DefaultError, QueryClient, QueryKey } from '@tanstack/query-core'
-import type { Injector } from '@angular/core'
+import { QueryObserver } from "@tanstack/query-core";
+import { runInInjectionContext } from "@angular/core";
+import { assertInjector } from "./util/assert-injector/assert-injector";
+import { injectQueryClient } from "./inject-query-client";
+import { createBaseQuery } from "./create-base-query";
+import { injectionContextProxy } from "./util/injection-context-proxy/injection-context-proxy";
+import type { DefaultError, QueryClient, QueryKey } from "@tanstack/query-core";
+import type { Injector } from "@angular/core";
 import type {
   CreateQueryOptions,
   CreateQueryResult,
   DefinedCreateQueryResult,
-} from './types'
+} from "./types";
 import type {
   DefinedInitialDataOptions,
   UndefinedInitialDataOptions,
-} from './query-options'
+} from "./query-options";
 
 export function injectQuery<
   TQueryFnData = unknown,
@@ -25,7 +26,7 @@ export function injectQuery<
     client: QueryClient,
   ) => UndefinedInitialDataOptions<TQueryFnData, TError, TData, TQueryKey>,
   injector?: Injector,
-): CreateQueryResult<TData, TError>
+): CreateQueryResult<TData, TError>;
 
 export function injectQuery<
   TQueryFnData = unknown,
@@ -37,7 +38,7 @@ export function injectQuery<
     client: QueryClient,
   ) => DefinedInitialDataOptions<TQueryFnData, TError, TData, TQueryKey>,
   injector?: Injector,
-): DefinedCreateQueryResult<TData, TError>
+): DefinedCreateQueryResult<TData, TError>;
 
 export function injectQuery<
   TQueryFnData = unknown,
@@ -49,20 +50,22 @@ export function injectQuery<
     client: QueryClient,
   ) => CreateQueryOptions<TQueryFnData, TError, TData, TQueryKey>,
   injector?: Injector,
-): CreateQueryResult<TData, TError>
+): CreateQueryResult<TData, TError>;
 
 export function injectQuery(
   options: (client: QueryClient) => CreateQueryOptions,
   injector?: Injector,
 ) {
-  const assertedInjector = assertInjector(injectQuery, injector)
+  const assertedInjector = assertInjector(injectQuery, injector);
   return assertInjector(injectQuery, injector, () => {
-    const queryClient = injectQueryClient()
+    const queryClient = injectQueryClient();
     return createBaseQuery(
       (client) =>
-        runInInjectionContext(assertedInjector, () => options(client)),
+        runInInjectionContext(assertedInjector, () =>
+          injectionContextProxy(assertedInjector, options(client)),
+        ),
       QueryObserver,
       queryClient,
-    )
-  })
+    );
+  });
 }
diff --git a/packages/angular-query-experimental/src/util/injection-context-proxy/injection-context-proxy.ts b/packages/angular-query-experimental/src/util/injection-context-proxy/injection-context-proxy.ts
new file mode 100644
index 0000000000..0421b7aa40
--- /dev/null
+++ b/packages/angular-query-experimental/src/util/injection-context-proxy/injection-context-proxy.ts
@@ -0,0 +1,37 @@
+import { runInInjectionContext } from "@angular/core";
+import type { Injector } from "@angular/core";
+
+export function injectionContextProxy<
+  TInput extends Record<string | symbol, any>,
+>(injector: Injector, input: TInput) {
+  return new Proxy<TInput>({} as TInput, {
+    get(target, prop) {
+      // first check if we have it in our internal state and return it
+      const computedField = target[prop];
+      if (computedField) {
+        return computedField;
+      }
+
+      const targetField = input[prop];
+      if (typeof targetField === "function") {
+        // @ts-expect-error
+        return (target[prop] = (...args: Array<any>) =>
+          runInInjectionContext(injector, () => targetField(...args)));
+      }
+
+      return targetField;
+    },
+    has(_, prop) {
+      return !!input[prop];
+    },
+    ownKeys() {
+      return Reflect.ownKeys(input);
+    },
+    getOwnPropertyDescriptor() {
+      return {
+        enumerable: true,
+        configurable: true,
+      };
+    },
+  });
+}