From a1a842e3d323b30154ecdac78e2847f7690089ce Mon Sep 17 00:00:00 2001
From: Bruno Lopes <bruno_lopes@live.com>
Date: Tue, 18 Jun 2024 14:49:12 -0300
Subject: [PATCH] chore: improve code repetition

---
 .../src/__tests__/prefetch.test.tsx           | 437 ++++++------------
 1 file changed, 137 insertions(+), 300 deletions(-)

diff --git a/packages/react-query/src/__tests__/prefetch.test.tsx b/packages/react-query/src/__tests__/prefetch.test.tsx
index 18c39b237e..63a946e68a 100644
--- a/packages/react-query/src/__tests__/prefetch.test.tsx
+++ b/packages/react-query/src/__tests__/prefetch.test.tsx
@@ -4,6 +4,8 @@ import { fireEvent, waitFor } from '@testing-library/react'
 import { ErrorBoundary } from 'react-error-boundary'
 import {
   QueryCache,
+  infiniteQueryOptions,
+  queryOptions,
   usePrefetchInfiniteQuery,
   usePrefetchQuery,
   useQueryErrorResetBoundary,
@@ -12,42 +14,83 @@ import {
 } from '..'
 import { createQueryClient, queryKey, renderWithClient, sleep } from './utils'
 
+import type { UseQueryOptions } from '..'
+import type { Mock } from 'vitest'
+
+function createQueryOptions<T>(params: { sleepInterval?: number; data: T }) {
+  const { sleepInterval = 10, data } = params
+
+  const queryFn = vi.fn().mockImplementation(async () => {
+    await sleep(sleepInterval)
+
+    return data
+  })
+
+  return queryOptions({
+    queryKey: queryKey(),
+    queryFn,
+  })
+}
+
+function createInfiniteQueryOptions(params: {
+  sleepInterval?: number
+  pages: Array<string>
+}) {
+  const { sleepInterval = 10, pages } = params
+
+  let currentPage = 0
+
+  return infiniteQueryOptions({
+    queryKey: queryKey(),
+    queryFn: vi.fn().mockImplementation(async () => {
+      const data = pages[currentPage] || undefined
+      await sleep(sleepInterval)
+      currentPage++
+
+      return {
+        page: currentPage,
+        totalPages: pages.length,
+        data,
+      }
+    }),
+    initialPageParam: 1,
+    getNextPageParam: (lastPage) => {
+      if (lastPage.page === lastPage.totalPages) {
+        return undefined
+      }
+
+      return lastPage.page + 1
+    },
+  })
+}
+
 describe('usePrefetchQuery', () => {
   const queryCache = new QueryCache()
   const queryClient = createQueryClient({ queryCache })
 
-  it('should prefetch query if query state does not exist', async () => {
-    const queryOpts = {
-      queryKey: queryKey(),
-      queryFn: vi.fn().mockImplementation(async () => {
-        await sleep(10)
-
-        return 'Prefetch is nice!'
-      }),
-    }
-
-    const suspendedQueryFn = vi.fn()
+  function Suspended(props: {
+    queryOpts: UseQueryOptions<any, Error, any, Array<string>>
+  }) {
+    const state = useSuspenseQuery(props.queryOpts)
 
-    function Suspended() {
-      const state = useSuspenseQuery({
-        ...queryOpts,
-        queryFn: suspendedQueryFn,
-      })
+    return (
+      <div>
+        <div>data: {String(state.data)}</div>
+      </div>
+    )
+  }
 
-      return (
-        <div>
-          <div>data: {String(state.data)}</div>
-          <div>fetching: {state.isFetching ? 'true' : 'false'}</div>
-        </div>
-      )
-    }
+  it('should prefetch query if query state does not exist', async () => {
+    const queryOpts = createQueryOptions({
+      data: 'Prefetch is nice!',
+    })
 
     function App() {
       usePrefetchQuery(queryOpts)
 
       return (
         <React.Suspense fallback="Loading...">
-          <Suspended />
+          <Suspended queryOpts={queryOpts} />
         </React.Suspense>
       )
     }
@@ -55,81 +98,38 @@ describe('usePrefetchQuery', () => {
     const rendered = renderWithClient(queryClient, <App />)
 
     await waitFor(() => rendered.getByText('data: Prefetch is nice!'))
-    expect(rendered.queryByText('fetching: true')).not.toBeInTheDocument()
     expect(queryOpts.queryFn).toHaveBeenCalledTimes(1)
-    expect(suspendedQueryFn).not.toHaveBeenCalled()
   })
 
   it('should not prefetch query if query state exists', async () => {
-    const queryFn = vi.fn().mockImplementation(async () => {
-      await sleep(10)
-
-      return 'Prefetch is really nice!'
-    })
-
-    const queryOpts = {
-      queryKey: queryKey(),
-      queryFn,
-    }
-
-    const suspendedQueryFn = vi.fn()
-
-    function Suspended() {
-      const state = useSuspenseQuery({
-        ...queryOpts,
-        queryFn: suspendedQueryFn,
-      })
-
-      return (
-        <div>
-          <div>data: {String(state.data)}</div>
-          <div>fetching: {state.isFetching ? 'true' : 'false'}</div>
-        </div>
-      )
-    }
+    const queryOpts = createQueryOptions({ data: 'Prefetch is really nice!' })
 
     function App() {
-      usePrefetchQuery({ ...queryOpts, queryFn: suspendedQueryFn })
+      usePrefetchQuery(queryOpts)
 
       return (
         <React.Suspense fallback="Loading...">
-          <Suspended />
+          <Suspended queryOpts={queryOpts} />
         </React.Suspense>
       )
     }
 
     await queryClient.fetchQuery(queryOpts)
-    queryFn.mockClear()
+    ;(queryOpts.queryFn as Mock).mockClear()
     const rendered = renderWithClient(queryClient, <App />)
 
     expect(rendered.queryByText('fetching: true')).not.toBeInTheDocument()
     await waitFor(() => rendered.getByText('data: Prefetch is really nice!'))
-    expect(suspendedQueryFn).not.toHaveBeenCalled()
     expect(queryOpts.queryFn).not.toHaveBeenCalled()
   })
 
   it('should display an error boundary if query cache is populated with an error', async () => {
-    const queryFn = vi.fn()
-
-    const queryOpts = {
-      queryKey: queryKey(),
-      queryFn,
-    }
-
-    const suspendedQueryFn = vi.fn()
-
-    function Suspended() {
-      const state = useSuspenseQuery({
-        ...queryOpts,
-        queryFn: suspendedQueryFn,
-      })
+    const queryOpts = createQueryOptions({ data: 'Not an error' })
+    ;(queryOpts.queryFn as Mock).mockImplementationOnce(async () => {
+      await sleep(10)
 
-      return (
-        <div>
-          <div>data: {String(state.data)}</div>
-        </div>
-      )
-    }
+      throw new Error('Oops! Server error!')
+    })
 
     function App() {
       usePrefetchQuery(queryOpts)
@@ -137,53 +137,28 @@ describe('usePrefetchQuery', () => {
       return (
         <ErrorBoundary fallbackRender={() => <div>Oops!</div>}>
           <React.Suspense fallback="Loading...">
-            <Suspended />
+            <Suspended queryOpts={queryOpts} />
           </React.Suspense>
         </ErrorBoundary>
       )
     }
 
-    queryFn.mockImplementationOnce(async () => {
-      await sleep(10)
-
-      throw new Error('Oops! Server error!')
-    })
-
     await queryClient.prefetchQuery(queryOpts)
-    queryFn.mockClear()
+    ;(queryOpts.queryFn as Mock).mockClear()
     const rendered = renderWithClient(queryClient, <App />)
 
     await waitFor(() => rendered.getByText('Oops!'))
-    expect(queryFn).not.toHaveBeenCalled()
-    expect(suspendedQueryFn).not.toHaveBeenCalled()
+    expect(rendered.queryByText('data: Not an error')).not.toBeInTheDocument()
+    expect(queryOpts.queryFn).not.toHaveBeenCalled()
   })
 
   it('should be able to recover from errors and try fetching again', async () => {
-    const queryFn = vi.fn()
-
-    const queryOpts = {
-      queryKey: queryKey(),
-      queryFn,
-    }
-
-    const suspendedQueryFn = vi.fn()
-
-    function Suspended() {
-      const state = useSuspenseQuery({
-        ...queryOpts,
-        queryFn: suspendedQueryFn.mockImplementationOnce(async () => {
-          await sleep(10)
-
-          return 'This is fine :dog: :fire:'
-        }),
-      })
+    const queryOpts = createQueryOptions({ data: 'This is fine :dog: :fire:' })
+    ;(queryOpts.queryFn as Mock).mockImplementationOnce(async () => {
+      await sleep(10)
 
-      return (
-        <div>
-          <div>data: {String(state.data)}</div>
-        </div>
-      )
-    }
+      throw new Error('Oops! Server error!')
+    })
 
     function App() {
       const { reset } = useQueryErrorResetBoundary()
@@ -200,96 +175,33 @@ describe('usePrefetchQuery', () => {
           )}
         >
           <React.Suspense fallback="Loading...">
-            <Suspended />
+            <Suspended queryOpts={queryOpts} />
           </React.Suspense>
         </ErrorBoundary>
       )
     }
 
-    queryFn.mockImplementationOnce(async () => {
-      await sleep(10)
-
-      throw new Error('Oops! Server error!')
-    })
-
     await queryClient.prefetchQuery(queryOpts)
+    ;(queryOpts.queryFn as Mock).mockClear()
 
     const rendered = renderWithClient(queryClient, <App />)
 
     await waitFor(() => rendered.getByText('Oops!'))
     fireEvent.click(rendered.getByText('Try again'))
     await waitFor(() => rendered.getByText('data: This is fine :dog: :fire:'))
-    expect(suspendedQueryFn).toHaveBeenCalled()
+    expect(queryOpts.queryFn).toHaveBeenCalledTimes(1)
   })
 
   it('should not create a suspense waterfall if prefetch is fired', async () => {
-    const firstQueryOpts = {
-      queryKey: queryKey(),
-      queryFn: vi.fn().mockImplementation(async () => {
-        await sleep(10)
-
-        return 'Prefetch is nice!'
-      }),
-    }
-
-    const secondQueryOpts = {
-      queryKey: queryKey(),
-      queryFn: vi.fn().mockImplementation(async () => {
-        await sleep(10)
-
-        return 'Prefetch is really nice!!'
-      }),
-    }
-
-    const thirdQueryOpts = {
-      queryKey: queryKey(),
-      queryFn: vi.fn().mockImplementation(async () => {
-        await sleep(10)
-
-        return 'Prefetch does not create waterfalls!!'
-      }),
-    }
+    const firstQueryOpts = createQueryOptions({ data: 'Prefetch is nice!' })
 
-    const suspendedQueryFn = vi.fn()
-
-    function FirstSuspended() {
-      const state = useSuspenseQuery({
-        ...firstQueryOpts,
-        queryFn: suspendedQueryFn,
-      })
-
-      return (
-        <div>
-          <div>data: {String(state.data)}</div>
-        </div>
-      )
-    }
-
-    function SecondSuspended() {
-      const state = useSuspenseQuery({
-        ...secondQueryOpts,
-        queryFn: suspendedQueryFn,
-      })
-
-      return (
-        <div>
-          <div>data: {String(state.data)}</div>
-        </div>
-      )
-    }
-
-    function ThirdSuspended() {
-      const state = useSuspenseQuery({
-        ...thirdQueryOpts,
-        queryFn: suspendedQueryFn,
-      })
+    const secondQueryOpts = createQueryOptions({
+      data: 'Prefetch is really nice!!',
+    })
 
-      return (
-        <div>
-          <div>data: {String(state.data)}</div>
-        </div>
-      )
-    }
+    const thirdQueryOpts = createQueryOptions({
+      data: 'Prefetch does not create waterfalls!!',
+    })
 
     const Fallback = vi.fn().mockImplementation(() => <div>Loading...</div>)
 
@@ -300,9 +212,9 @@ describe('usePrefetchQuery', () => {
 
       return (
         <React.Suspense fallback={<Fallback />}>
-          <FirstSuspended />
-          <SecondSuspended />
-          <ThirdSuspended />
+          <Suspended queryOpts={firstQueryOpts} />
+          <Suspended queryOpts={secondQueryOpts} />
+          <Suspended queryOpts={thirdQueryOpts} />
         </React.Suspense>
       )
     }
@@ -316,7 +228,9 @@ describe('usePrefetchQuery', () => {
       rendered.getByText('data: Prefetch does not create waterfalls!!'),
     )
     expect(Fallback).toHaveBeenCalledTimes(1)
-    expect(suspendedQueryFn).not.toHaveBeenCalled()
+    expect(firstQueryOpts.queryFn).toHaveBeenCalledTimes(1)
+    expect(secondQueryOpts.queryFn).toHaveBeenCalledTimes(1)
+    expect(thirdQueryOpts.queryFn).toHaveBeenCalledTimes(1)
   })
 })
 
@@ -324,160 +238,83 @@ describe('usePrefetchInfiniteQuery', () => {
   const queryCache = new QueryCache()
   const queryClient = createQueryClient({ queryCache })
 
-  it('should prefetch an infinite query if query state does not exist', async () => {
-    const pages = [
-      'Do you fetch on render?',
-      'Or do you render as you fetch?',
-      'Either way, Tanstack Query helps you!',
-    ]
-
-    let queryPage = 0
-
-    const queryOpts = {
-      queryKey: queryKey(),
-      queryFn: vi
-        .fn<any, Promise<{ page: number; totalPages: number; data: string }>>()
-        .mockImplementation(async () => {
-          const data = pages[queryPage] || ''
-          await sleep(10)
-          queryPage++
-          return {
-            page: queryPage,
-            totalPages: pages.length,
-            data,
-          }
-        }),
-      initialPageParam: 1,
-      getNextPageParam: (lastPage: any) => {
-        if (lastPage.page === lastPage.totalPages) {
-          return undefined
-        }
-
-        return lastPage.page + 1
-      },
-      pages: 3,
-    }
+  const Fallback = vi.fn().mockImplementation(() => <div>Loading...</div>)
 
-    const suspendedQueryFn = vi.fn()
+  function Suspended(props: { queryOpts: any }) {
+    const state = useSuspenseInfiniteQuery(props.queryOpts)
 
-    function Suspended() {
-      const state = useSuspenseInfiniteQuery({
-        ...queryOpts,
-        queryFn: suspendedQueryFn,
-      })
-
-      return (
-        <div>
-          {state.data.pages.map((page) => (
+    return (
+      <div>
+        {(state.data.pages as Array<{ page: number; data: string }>).map(
+          (page) => (
             <div key={page.page}>data: {page.data}</div>
-          ))}
-          <button onClick={() => state.fetchNextPage()}>Next Page</button>
-        </div>
-      )
-    }
+          ),
+        )}
+        <button onClick={() => state.fetchNextPage()}>Next Page</button>
+      </div>
+    )
+  }
+
+  it('should prefetch an infinite query if query state does not exist', async () => {
+    const queryOpts = createInfiniteQueryOptions({
+      pages: [
+        'Do you fetch on render?',
+        'Or do you render as you fetch?',
+        'Either way, Tanstack Query helps you!',
+      ],
+    })
 
     function App() {
       usePrefetchInfiniteQuery(queryOpts)
 
       return (
-        <React.Suspense fallback="Loading...">
-          <Suspended />
+        <React.Suspense fallback={<Fallback />}>
+          <Suspended queryOpts={queryOpts} />
         </React.Suspense>
       )
     }
 
     const rendered = renderWithClient(queryClient, <App />)
 
-    await waitFor(() => rendered.getByText('Loading...'))
     await waitFor(() => rendered.getByText('data: Do you fetch on render?'))
     fireEvent.click(rendered.getByText('Next Page'))
-    expect(rendered.queryByText('Loading...')).not.toBeInTheDocument()
     await waitFor(() =>
       rendered.getByText('data: Or do you render as you fetch?'),
     )
     fireEvent.click(rendered.getByText('Next Page'))
-    expect(rendered.queryByText('Loading...')).not.toBeInTheDocument()
     await waitFor(() =>
       rendered.getByText('data: Either way, Tanstack Query helps you!'),
     )
+    expect(Fallback).toHaveBeenCalledTimes(1)
     expect(queryOpts.queryFn).toHaveBeenCalledTimes(3)
-    expect(suspendedQueryFn).not.toHaveBeenCalled()
   })
 
-  it('should not prefetch an infinite query if query state already exists', async () => {
-    const pages = [
-      'Prefetch rocks!',
-      'No waterfalls, boy!',
-      'Tanstack Query #ftw',
-    ]
-
-    let queryPage = 0
-
-    const suspendedQueryFn = vi.fn()
-
-    const queryOpts = {
-      queryKey: queryKey(),
-      queryFn: vi
-        .fn<any, Promise<{ page: number; totalPages: number; data: string }>>()
-        .mockImplementation(async () => {
-          const data = pages[queryPage] || ''
-          await sleep(10)
-          queryPage++
-          return {
-            page: queryPage,
-            totalPages: pages.length,
-            data,
-          }
-        }),
-      initialPageParam: 1,
-      getNextPageParam: (lastPage: any) => {
-        if (lastPage.page === lastPage.totalPages) {
-          return undefined
-        }
-
-        return lastPage.page + 1
-      },
-      pages: 3,
-    }
-
-    function Suspended() {
-      const state = useSuspenseInfiniteQuery({
-        ...queryOpts,
-        queryFn: suspendedQueryFn,
-      })
-
-      return (
-        <div>
-          {state.data.pages.map((page) => (
-            <div key={page.page}>data: {page.data}</div>
-          ))}
-          <button onClick={() => state.fetchNextPage()}>Next Page</button>
-        </div>
-      )
-    }
+  it('should not display fallback if the query cache is already populated', async () => {
+    const queryOpts = createInfiniteQueryOptions({
+      pages: ['Prefetch rocks!', 'No waterfalls, boy!', 'Tanstack Query #ftw'],
+    })
 
-    await queryClient.prefetchInfiniteQuery(queryOpts)
+    await queryClient.prefetchInfiniteQuery({ ...queryOpts, pages: 3 })
+    ;(queryOpts.queryFn as Mock).mockClear()
 
     function App() {
-      usePrefetchInfiniteQuery({ ...queryOpts, queryFn: suspendedQueryFn })
+      usePrefetchInfiniteQuery(queryOpts)
 
       return (
-        <React.Suspense fallback="Loading...">
-          <Suspended />
+        <React.Suspense fallback={<Fallback />}>
+          <Suspended queryOpts={queryOpts} />
         </React.Suspense>
       )
     }
 
     const rendered = renderWithClient(queryClient, <App />)
 
-    expect(rendered.queryByText('Loading...')).not.toBeInTheDocument()
     await waitFor(() => rendered.getByText('data: Prefetch rocks!'))
     fireEvent.click(rendered.getByText('Next Page'))
-    expect(rendered.queryByText('Loading...')).not.toBeInTheDocument()
     await waitFor(() => rendered.getByText('data: No waterfalls, boy!'))
     fireEvent.click(rendered.getByText('Next Page'))
-    expect(rendered.queryByText('Loading...')).not.toBeInTheDocument()
     await waitFor(() => rendered.getByText('data: Tanstack Query #ftw'))
-    expect(suspendedQueryFn).not.toHaveBeenCalled()
+    expect(queryOpts.queryFn).not.toHaveBeenCalled()
+    expect(Fallback).not.toHaveBeenCalled()
   })
 })