import type { QueryKey } from '@tanstack/react-query' interface QueryFilter { queryKey: QueryKey exact?: boolean } type Updater = T | ((previous: T | undefined) => T | undefined) interface StoredQueryEntry { queryKey: QueryKey data: unknown } function isPrefixQueryKey(target: QueryKey, prefix: QueryKey): boolean { if (prefix.length > target.length) return false return prefix.every((value, index) => Object.is(value, target[index])) } function keyOf(queryKey: QueryKey): string { return JSON.stringify(queryKey) } export class MockQueryClient { private readonly queryMap = new Map() async cancelQueries(filters: QueryFilter): Promise { void filters } seedQuery(queryKey: QueryKey, data: T | undefined) { this.queryMap.set(keyOf(queryKey), { queryKey, data, }) } getQueryData(queryKey: QueryKey): T | undefined { const entry = this.queryMap.get(keyOf(queryKey)) return entry?.data as T | undefined } setQueryData(queryKey: QueryKey, updater: Updater) { const previous = this.getQueryData(queryKey) const next = typeof updater === 'function' ? (updater as (prev: T | undefined) => T | undefined)(previous) : updater this.seedQuery(queryKey, next) } getQueriesData(filters: QueryFilter): Array<[QueryKey, T | undefined]> { const matched: Array<[QueryKey, T | undefined]> = [] for (const { queryKey, data } of this.queryMap.values()) { const isMatch = filters.exact ? keyOf(filters.queryKey) === keyOf(queryKey) : isPrefixQueryKey(queryKey, filters.queryKey) if (!isMatch) continue matched.push([queryKey, data as T | undefined]) } return matched } setQueriesData( filters: QueryFilter, updater: (previous: T | undefined) => T | undefined, ) { const matches = this.getQueriesData(filters) matches.forEach(([queryKey, previous]) => { this.seedQuery(queryKey, updater(previous)) }) } }