refactor: add more tests

src/define-query.spec.ts

@@ -1,18 +1,19 @@
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
-import { enableAutoUnmount, mount } from '@vue/test-utils'
+import { enableAutoUnmount, flushPromises, mount } from '@vue/test-utils'
 import { createPinia } from 'pinia'
-import { ref } from 'vue'
+import type { App } from 'vue'
+import { createApp, defineComponent, ref } from 'vue'
 import { QueryPlugin } from './query-plugin'
 import { defineQuery } from './define-query'
 import { useQuery } from './use-query'
 
 describe('defineQuery', () => {
-    beforeEach(() => {
-      vi.useFakeTimers()
-    })
-    afterEach(() => {
-      vi.restoreAllMocks()
-    })
+  beforeEach(() => {
+    vi.useFakeTimers()
+  })
+  afterEach(() => {
+    vi.restoreAllMocks()
+  })
 
   enableAutoUnmount(afterEach)
 
@@ -75,4 +76,224 @@ describe('defineQuery', () => {
     expect(todoFilter).toBe(useTodoList().todoFilter)
     expect(todoFilter).toBe(returnedValues.todoFilter)
   })
+
+  describe('refetchOnMount', () => {
+    it('refreshes the query if mounted in a new component', async () => {
+      const spy = vi.fn(async () => {
+        return 'todos'
+      })
+      const useTodoList = defineQuery({
+        key: ['todos'],
+        query: spy,
+        refetchOnMount: true,
+        staleTime: 100,
+      })
+      let returnedValues!: ReturnType<typeof useTodoList>
+
+      const pinia = createPinia()
+      const Component = defineComponent({
+        setup() {
+          returnedValues = useTodoList()
+          return {}
+        },
+        template: `<div></div>`,
+      })
+
+      mount(Component, {
+        global: {
+          plugins: [pinia, QueryPlugin],
+        },
+      })
+      await flushPromises()
+
+      const { data, status } = returnedValues
+      expect(spy).toHaveBeenCalledTimes(1)
+      expect(status.value).toBe('success')
+      expect(data.value).toEqual('todos')
+
+      mount(Component, {
+        global: {
+          plugins: [pinia, QueryPlugin],
+        },
+      })
+      // still called only once
+      expect(spy).toHaveBeenCalledTimes(1)
+      // no ongoing call
+      expect(status.value).toBe('success')
+      expect(data.value).toEqual('todos')
+      await flushPromises()
+      expect(status.value).toBe('success')
+
+      vi.advanceTimersByTime(101)
+      mount(Component, {
+        global: {
+          plugins: [pinia, QueryPlugin],
+        },
+      })
+      // it should be loading
+      expect(status.value).toBe('loading')
+      expect(data.value).toEqual('todos')
+      await flushPromises()
+
+      expect(spy).toHaveBeenCalledTimes(2)
+    })
+
+    it('refetches if refetchOnMount is always', async () => {
+      const spy = vi.fn(async () => {
+        return 'todos'
+      })
+      const useTodoList = defineQuery({
+        key: ['todos'],
+        query: spy,
+        refetchOnMount: 'always',
+        staleTime: 100,
+      })
+      let returnedValues!: ReturnType<typeof useTodoList>
+
+      const pinia = createPinia()
+      const Component = defineComponent({
+        setup() {
+          returnedValues = useTodoList()
+          return {}
+        },
+        template: `<div></div>`,
+      })
+
+      mount(Component, {
+        global: {
+          plugins: [pinia, QueryPlugin],
+        },
+      })
+      await flushPromises()
+
+      const { data, status } = returnedValues
+      expect(spy).toHaveBeenCalledTimes(1)
+      expect(status.value).toBe('success')
+      expect(data.value).toEqual('todos')
+
+      mount(Component, {
+        global: {
+          plugins: [pinia, QueryPlugin],
+        },
+      })
+      expect(spy).toHaveBeenCalledTimes(2)
+    })
+
+    it('does not refetch if refetchOnMount is false', async () => {
+      const spy = vi.fn(async () => {
+        return 'todos'
+      })
+      const useTodoList = defineQuery({
+        key: ['todos'],
+        query: spy,
+        refetchOnMount: false,
+        staleTime: 100,
+      })
+      let returnedValues!: ReturnType<typeof useTodoList>
+
+      const pinia = createPinia()
+      const Component = defineComponent({
+        setup() {
+          returnedValues = useTodoList()
+          return {}
+        },
+        template: `<div></div>`,
+      })
+
+      mount(Component, {
+        global: {
+          plugins: [pinia, QueryPlugin],
+        },
+      })
+      await flushPromises()
+
+      const { data, status } = returnedValues
+      expect(spy).toHaveBeenCalledTimes(1)
+      expect(status.value).toBe('success')
+      expect(data.value).toEqual('todos')
+
+      mount(Component, {
+        global: {
+          plugins: [pinia, QueryPlugin],
+        },
+      })
+      expect(spy).toHaveBeenCalledTimes(1)
+    })
+  })
+
+  describe('outside of components', () => {
+    let app: App
+    beforeEach(() => {
+      const pinia = createPinia()
+      app = createApp({ render: () => null })
+        .use(pinia)
+        .use(QueryPlugin)
+      app.mount(document.createElement('div'))
+    })
+    afterEach(() => {
+      app?.unmount()
+    })
+
+    it('reuses the query', async () => {
+      const useTodoList = defineQuery({
+        key: ['todos'],
+        query: async () => [{ id: 1 }],
+      })
+
+      // to have access to inject
+      app.runWithContext(() => {
+        const { data } = useTodoList()
+        expect(data).toBe(useTodoList().data)
+      })
+    })
+
+    it('reuses the query with a setup function', async () => {
+      const useTodoList = defineQuery(() => {
+        const todoFilter = ref<'all' | 'finished' | 'unfinished'>('all')
+        const { data, ...rest } = useQuery({
+          key: ['todos', { filter: todoFilter.value }],
+          query: async () => [{ id: 1 }],
+        })
+        return { ...rest, todoList: data, todoFilter }
+      })
+
+      app.runWithContext(() => {
+        const { todoList, todoFilter } = useTodoList()
+        expect(todoList).toBe(useTodoList().todoList)
+        expect(todoFilter).toBe(useTodoList().todoFilter)
+      })
+    })
+
+    it('refreshes the query when called', async () => {
+      const spy = vi.fn(async () => {
+        return 'todos'
+      })
+      const useTodoList = defineQuery({
+        key: ['todos'],
+        query: spy,
+        staleTime: 100,
+        refetchOnMount: true,
+      })
+
+      await app.runWithContext(async () => {
+        const { data, status } = useTodoList()
+        await flushPromises()
+        expect(spy).toHaveBeenCalledTimes(1)
+
+        expect(status.value).toBe('success')
+        expect(data.value).toEqual('todos')
+
+        // should not trigger a refresh
+        useTodoList()
+        await flushPromises()
+        expect(spy).toHaveBeenCalledTimes(1)
+
+        vi.advanceTimersByTime(101)
+
+        useTodoList()
+        await flushPromises()
+        expect(spy).toHaveBeenCalledTimes(2)
+      })
+    })
+  })
 })

src/define-query.ts

@@ -17,16 +17,14 @@ import { useQuery } from './use-query'
  * })
  * ```
  */
-export function defineQuery<
-  TResult,
-  TError = ErrorDefault,
->(
+export function defineQuery<TResult, TError = ErrorDefault>(
   options: UseQueryOptions<TResult, TError>,
 ): () => UseQueryReturn<TResult, TError>
 
 /**
  * Define a query with a setup function. Allows to return arbitrary values from the query function, create contextual
- * refs, rename the returned values, etc.
+ * refs, rename the returned values, etc. The setup function will be called only once, like stores, and **must be
+ * synchronous**.
  *
  * @param setup - a function to setup the query
  * @example
@@ -47,6 +45,9 @@ export function defineQuery<T>(setup: () => T): () => T
 export function defineQuery(
   optionsOrSetup: UseQueryOptions | (() => unknown),
 ): () => unknown {
-  const setupFn = typeof optionsOrSetup === 'function' ? optionsOrSetup : () => useQuery(optionsOrSetup)
+  const setupFn
+    = typeof optionsOrSetup === 'function'
+      ? optionsOrSetup
+      : () => useQuery(optionsOrSetup)
   return () => useQueryCache().ensureDefinedQuery(setupFn)
 }

src/query-store.ts

@@ -163,25 +163,39 @@ export const useQueryCache = defineStore(QUERY_STORE_ID, () => {
   // this allows use to attach reactive effects to the scope later on
   const scope = getCurrentScope()!
 
-  const defineQueryMap = new WeakMap<() => unknown, { queries: string[][], result: any }>()
-  let currentDefineQuerySetupFunction: (() => unknown) | null
+  type DefineQueryEntry = [entries: UseQueryEntry[], returnValue: unknown]
+  // keep track of the entry being defined so we can add the queries in ensureEntry
+  // this allows us to refresh the entry when a defined query is used again
+  // and refetchOnMount is true
+  let currentDefineQueryEntry: DefineQueryEntry | undefined | null
+  const defineQueryMap = new WeakMap<() => unknown, DefineQueryEntry>()
   function ensureDefinedQuery<T>(fn: () => T): T {
-    if (!defineQueryMap.has(fn)) {
-      currentDefineQuerySetupFunction = fn
-      defineQueryMap.set(fn, { queries: [], result: null })
-      defineQueryMap.get(fn)!.result = scope.run(fn)!
-      currentDefineQuerySetupFunction = null
+    let defineQueryEntry = defineQueryMap.get(fn)
+    if (!defineQueryEntry) {
+      // create the entry first
+      currentDefineQueryEntry = defineQueryEntry = [[], null]
+      // then run it s oit can add the queries to the entry
+      defineQueryEntry[1] = scope.run(fn)
+      currentDefineQueryEntry = null
+      defineQueryMap.set(fn, defineQueryEntry)
     } else {
-      defineQueryMap.get(fn)!.queries.forEach(
-        (key) => {
-          const query = cachesRaw.get(key) as UseQueryEntry | undefined
-          if (query) refresh(query)
-        },
-      )
+      // if the entry already exists, we know the queries inside
+      for (const queryEntry of defineQueryEntry[0]) {
+        // TODO: refactor this to be a method of the store so it can be used in useQuery too
+        // and not be called during hydration
+        if (queryEntry.options?.refetchOnMount) {
+          if (queryEntry.options.refetchOnMount === 'always') {
+            refetch(queryEntry)
+          } else {
+            // console.log('refreshing')
+            refresh(queryEntry)
+          }
+        }
+      }
     }
-    return defineQueryMap.get(fn)!.result
-  }
 
+    return defineQueryEntry[1] as T
+  }
   function ensureEntry<TResult = unknown, TError = ErrorDefault>(
     keyRaw: EntryKey,
     options: UseQueryOptionsWithDefaults<TResult, TError>,
@@ -191,11 +205,8 @@ export const useQueryCache = defineStore(QUERY_STORE_ID, () => {
         `useQuery() was called with an empty array as the key. It must have at least one element.`,
       )
     }
+
     const key = keyRaw.map(stringifyFlatObject)
-    if (currentDefineQuerySetupFunction) {
-      const currentDefineQueryEntry = defineQueryMap.get(currentDefineQuerySetupFunction)
-      currentDefineQueryEntry!.queries.push(key)
-    }
     // ensure the state
     // console.log('⚙️ Ensuring entry', key)
     let entry = cachesRaw.get(key) as UseQueryEntry<TResult, TError> | undefined
@@ -208,8 +219,12 @@ export const useQueryCache = defineStore(QUERY_STORE_ID, () => {
       )
     }
 
-    if (!entry.options) {
-      entry.options = options
+    // add the options to the entry the first time only
+    entry.options ??= options
+
+    // if this query was defined within a defineQuery call, add it to the list
+    if (currentDefineQueryEntry) {
+      currentDefineQueryEntry[0].push(entry)
     }
 
     return entry