From 40bd67a17388b8c986b6a9fece93f06e8a4b3e50 Mon Sep 17 00:00:00 2001
From: Arnoud de Vries <6420061+arnoud-dv@users.noreply.github.com>
Date: Sat, 23 Dec 2023 02:40:31 +0100
Subject: [PATCH] fix(angular-query): fix optimistic result

---
 .../src/createBaseQuery.ts                    | 52 +++++++++----------
 .../src/tests/injectQuery.test.ts             | 16 +-----
 2 files changed, 26 insertions(+), 42 deletions(-)

diff --git a/packages/angular-query-experimental/src/createBaseQuery.ts b/packages/angular-query-experimental/src/createBaseQuery.ts
index a6c789e7f6..fa739a494f 100644
--- a/packages/angular-query-experimental/src/createBaseQuery.ts
+++ b/packages/angular-query-experimental/src/createBaseQuery.ts
@@ -6,19 +6,12 @@ import {
   inject,
   signal,
 } from '@angular/core'
-import { notifyManager } from '@tanstack/query-core'
 import { createResultStateSignalProxy } from './query-proxy'
-import type { Signal } from '@angular/core'
-import type {
-  QueryClient,
-  QueryKey,
-  QueryObserver,
-  QueryObserverResult,
-} from '@tanstack/query-core'
+import type { QueryClient, QueryKey, QueryObserver } from '@tanstack/query-core'
 import type { CreateBaseQueryOptions, CreateBaseQueryResult } from './types'
 
 /**
- * Base implementation for `createQuery` and `createInfiniteQuery`.
+ * Base implementation for `injectQuery` and `injectInfiniteQuery`.
  * @internal
  */
 export function createBaseQuery<
@@ -43,7 +36,12 @@ export function createBaseQuery<
   assertInInjectionContext(createBaseQuery)
   const destroyRef = inject(DestroyRef)
 
-  /** Creates a signal that has the default options applied */
+  /**
+   * Signal that has the default options from query client applied
+   * computed() is used so signals can be inserted into the options
+   * making it reactive. Wrapping options in a function ensures embedded expressions
+   * are preserved and can keep being applied after signal changes
+   */
   const defaultedOptionsSignal = computed(() => {
     const defaultedOptions = queryClient.defaultQueryOptions(
       options(queryClient),
@@ -52,7 +50,6 @@ export function createBaseQuery<
     return defaultedOptions
   })
 
-  /** Creates the observer */
   const observer = new Observer<
     TQueryFnData,
     TError,
@@ -61,27 +58,28 @@ export function createBaseQuery<
     TQueryKey
   >(queryClient, defaultedOptionsSignal())
 
-  effect(() => {
-    // Do not notify on updates because of changes in the options because
-    // these changes should already be reflected in the optimistic result.
-    observer.setOptions(defaultedOptionsSignal(), { listeners: false })
-  })
-
-  const result = signal(observer.getCurrentResult())
-
-  const unsubscribe = observer.subscribe(
-    notifyManager.batchCalls((val) => result.set(val)),
+  const resultSignal = signal(
+    observer.getOptimisticResult(
+      queryClient.defaultQueryOptions(defaultedOptionsSignal()),
+    ),
   )
-  destroyRef.onDestroy(unsubscribe)
 
-  /** Subscribe to changes in result and defaultedOptionsStore */
-  const resultSignal: Signal<QueryObserverResult<TData, TError>> = computed(
+  effect(
     () => {
-      return !defaultedOptionsSignal().notifyOnChangeProps
-        ? observer.trackResult(result())
-        : result()
+      // Do not notify on updates because of changes in the options because
+      // these changes should already be reflected in the optimistic result.
+      const defaultedOptions = defaultedOptionsSignal()
+      observer.setOptions(defaultedOptions, {
+        listeners: false,
+      })
+      resultSignal.set(observer.getOptimisticResult(defaultedOptions))
     },
+    { allowSignalWrites: true },
   )
 
+  // observer.trackResult is not used as this optimization is not needed for Angular
+  const unsubscribe = observer.subscribe(resultSignal.set)
+  destroyRef.onDestroy(unsubscribe)
+
   return createResultStateSignalProxy<TData, TError>(resultSignal)
 }
diff --git a/packages/angular-query-experimental/src/tests/injectQuery.test.ts b/packages/angular-query-experimental/src/tests/injectQuery.test.ts
index 5aba333de3..bb0a399649 100644
--- a/packages/angular-query-experimental/src/tests/injectQuery.test.ts
+++ b/packages/angular-query-experimental/src/tests/injectQuery.test.ts
@@ -147,7 +147,7 @@ describe('injectQuery', () => {
     TestBed.flushEffects()
 
     expect(query1.data()).toStrictEqual('Some data')
-    // expect(query2().fetchStatus).toStrictEqual('fetching') // TODO: is this working correctly?
+    expect(query2.fetchStatus()).toStrictEqual('fetching')
 
     flush()
 
@@ -215,18 +215,4 @@ describe('injectQuery', () => {
 
     expect(query.status()).toBe('error')
   }))
-
-  it('should not update signal when notifyOnChangeProps is set without the changed property being in notifyOnChangeProps', fakeAsync(() => {
-    const query = TestBed.runInInjectionContext(() => {
-      return injectQuery(() => ({
-        queryKey: ['key14'],
-        queryFn: simpleFetcher,
-        notifyOnChangeProps: 'all',
-      }))
-    })
-
-    flush()
-
-    expect(query.status()).toBe('success')
-  }))
 })