fix(react-query): non continuous suspense with useSuspenseQueries (#6298) (#6303)

Co-authored-by: Dominik Dorfmeister <[email protected]>

Author: maciej-baruch

packages/react-query/src/__tests__/useSuspenseQueries.test.tsx

@@ -0,0 +1,144 @@
+import {
+  afterAll,
+  afterEach,
+  beforeAll,
+  describe,
+  expect,
+  it,
+  vi,
+} from 'vitest'
+import { act, render } from '@testing-library/react'
+import React from 'react'
+import { useSuspenseQueries } from '..'
+import { createQueryClient, sleep } from './utils'
+import type { UseSuspenseQueryOptions } from '..'
+
+type NumberQueryOptions = UseSuspenseQueryOptions<number>
+
+const QUERY_DURATION = 1000
+
+const createQuery: (id: number) => NumberQueryOptions = (id) => ({
+  queryKey: [id],
+  queryFn: async () => {
+    await sleep(QUERY_DURATION)
+    return id
+  },
+})
+const resolveQueries = () => vi.advanceTimersByTimeAsync(QUERY_DURATION)
+
+const queryClient = createQueryClient()
+
+describe('useSuspenseQueries', () => {
+  const onSuspend = vi.fn()
+  const onQueriesResolution = vi.fn()
+
+  beforeAll(() => {
+    vi.useFakeTimers()
+  })
+
+  afterAll(() => {
+    vi.useRealTimers()
+  })
+
+  afterEach(() => {
+    queryClient.clear()
+    onSuspend.mockClear()
+    onQueriesResolution.mockClear()
+  })
+
+  function SuspenseFallback() {
+    React.useEffect(() => {
+      onSuspend()
+    }, [])
+
+    return null
+  }
+
+  const withSuspenseWrapper = <T extends object>(Component: React.FC<T>) => {
+    function SuspendedComponent(props: T) {
+      return (
+        <React.Suspense fallback={<SuspenseFallback />}>
+          <Component {...props} />
+        </React.Suspense>
+      )
+    }
+
+    return SuspendedComponent
+  }
+
+  function QueriesContainer({
+    queries,
+  }: {
+    queries: Array<NumberQueryOptions>
+  }) {
+    const queriesResults = useSuspenseQueries(
+      { queries, combine: (results) => results.map((r) => r.data) },
+      queryClient,
+    )
+
+    React.useEffect(() => {
+      onQueriesResolution(queriesResults)
+    }, [queriesResults])
+
+    return null
+  }
+
+  const TestComponent = withSuspenseWrapper(QueriesContainer)
+
+  it('should suspend on mount', () => {
+    render(<TestComponent queries={[1, 2].map(createQuery)} />)
+
+    expect(onSuspend).toHaveBeenCalledOnce()
+  })
+
+  it('should resolve queries', async () => {
+    render(<TestComponent queries={[1, 2].map(createQuery)} />)
+
+    await act(resolveQueries)
+
+    expect(onQueriesResolution).toHaveBeenCalledTimes(1)
+    expect(onQueriesResolution).toHaveBeenLastCalledWith([1, 2])
+  })
+
+  it('should not suspend on mount if query has been already fetched', () => {
+    const query = createQuery(1)
+
+    queryClient.setQueryData(query.queryKey, query.queryFn)
+
+    render(<TestComponent queries={[query]} />)
+
+    expect(onSuspend).not.toHaveBeenCalled()
+  })
+
+  it('should not break suspense when queries change without resolving', async () => {
+    const initQueries = [1, 2].map(createQuery)
+    const nextQueries = [3, 4, 5, 6].map(createQuery)
+
+    const { rerender } = render(<TestComponent queries={initQueries} />)
+
+    rerender(<TestComponent queries={nextQueries} />)
+
+    await act(resolveQueries)
+
+    expect(onSuspend).toHaveBeenCalledTimes(1)
+    expect(onQueriesResolution).toHaveBeenCalledTimes(1)
+    expect(onQueriesResolution).toHaveBeenLastCalledWith([3, 4, 5, 6])
+  })
+
+  it('should suspend only once per queries change', async () => {
+    const initQueries = [1, 2].map(createQuery)
+    const nextQueries = [3, 4, 5, 6].map(createQuery)
+
+    const { rerender } = render(<TestComponent queries={initQueries} />)
+
+    await act(resolveQueries)
+
+    rerender(<TestComponent queries={nextQueries} />)
+
+    await act(resolveQueries)
+
+    expect(onSuspend).toHaveBeenCalledTimes(2)
+    expect(onQueriesResolution).toHaveBeenCalledTimes(2)
+    expect(onQueriesResolution).toHaveBeenLastCalledWith([3, 4, 5, 6])
+  })
+})

packages/react-query/src/useQueries.ts

@@ -1,7 +1,11 @@
 'use client'
 import * as React from 'react'
 
-import { QueriesObserver, notifyManager } from '@tanstack/query-core'
+import {
+  QueriesObserver,
+  QueryObserver,
+  notifyManager,
+} from '@tanstack/query-core'
 import { useQueryClient } from './QueryClientProvider'
 import { useIsRestoring } from './isRestoring'
 import { useQueryErrorResetBoundary } from './QueryErrorResetBoundary'
@@ -262,9 +266,9 @@ export function useQueries<
   const suspensePromises = shouldAtLeastOneSuspend
     ? optimisticResult.flatMap((result, index) => {
         const opts = defaultedQueries[index]
-        const queryObserver = observer.getObservers()[index]
 
-        if (opts && queryObserver) {
+        if (opts) {
+          const queryObserver = new QueryObserver(client, opts)
           if (shouldSuspend(opts, result, isRestoring)) {
             return fetchOptimistic(opts, queryObserver, errorResetBoundary)
           } else if (willFetch(result, isRestoring)) {