Storybook - React - tanstack query mock. An approach to mock states and data for useQuery.

Component.stories.tsx

import { Meta, StoryFn } from "@storybook/react";
import { QueryClientProvider } from "@tanstack/react-query";
import {
  mockQuery,
  mockQueryError,
  mockQueryLoading,
} from "@/lib/mocks/mock-query";
import { Component } from "./Component";


const meta = {
  component: Component,
} satisfies Meta<typeof Component>;

export default meta;

type Story = StoryFn<{
  state: "loading" | "empty" | "data" | "error";
  data?: any[];
}>;

const Template: Story = (args) => {
  let queryClient;

  switch (args.state) {
    case "loading":
      queryClient = mockQueryLoading(["yourKey"]);
      break;
    case "empty":
      queryClient = mockQuery(["yourKey"], []);
      break;
    case "data": 
      queryClient = mockQuery(["yourKey"], args.data);
      break;
    case "error":
      queryClient = mockQueryError(["yourKey"], new Error("Failed to load roles"));
      break;
    default:
      queryClient = mockQueryLoading(["yourKey"]);
  }

  return (
    <QueryClientProvider client={queryClient}>
      <Component />;
    </QueryClientProvider>
  );
};

export const Loading: Story = Template.bind({});

export const Empty: Story = Template.bind({});
Empty.args = {
  state: "empty",
};

export const ErrorState = Template.bind({});
ErrorState.args = {
  state: "error",
};

export const Data: Story = Template.bind({});
Data.args = {
  state: "data",
  data: [
    {
      id: "1",
      name: "Data 1",
    },
    {
      id: "2",
      name: "Data 2",
    },
  ],
};

Component.tsx

import { useGetData } from "./useGetData";

export const Component = (): JSX.Element => {

    const { isLoading, isError, data } = useGetRoles();

    if (isLoading) {
        return (
            <div>loading...</div>
        );
    }

    if (isError) {
        return (
            <div>Error on fetching data!</div>
        );
    }

    if (data?.length === 0) {
        return (
            <div>No data!</div>
        );
    }

    return (
      <div>
        {/* Render you data here */}
      </div>
    );
};

mock-query.ts

import {
  DefaultError,
  Query,
  QueryCache,
  QueryClient,
  QueryKey,
  QueryState,
} from "@tanstack/react-query";

class QueryMock extends Query {
  constructor(private key: QueryKey, public state: QueryState) {
    super({
      queryKey: key,
      options: {
        queryFn: async () => {
          if (state.error) {
            throw state.error;
          }
          return state.data;
        },
        retry: false,
        
      },
      defaultOptions: {
        queryFn: async () => {
          if (state.error) {
            throw state.error;
          }
          return state.data;
        },
        retry: false,
      },
      queryHash: "hash",
      cache: new QueryCache(),
      state,
      
    });
  }
}

class QueryCacheMock extends QueryCache {

  constructor(private queryKey: QueryKey, private state: QueryState) {
    super();
  }

  get<
    TQueryFnData = unknown,
    TError = DefaultError,
    TData = TQueryFnData,
    TQueryKey extends QueryKey = QueryKey,
  >(
    queryHash: string,
  ): Query<TQueryFnData, TError, TData, TQueryKey> | undefined  {
    const queryKey = JSON.parse(queryHash) as QueryKey;

    if (JSON.stringify(queryKey) !== JSON.stringify(this.queryKey)) {
      return undefined;
    }

    return new QueryMock(this.queryKey, this.state) as unknown as Query<TQueryFnData, TError, TData, TQueryKey>;
  }
}

class QueryClientMock extends QueryClient {

  private readonly queryCache: QueryCacheMock;

  constructor(
    private queryKey: QueryKey,
    private queryState: QueryState<any, any>
  ) {
    super({
      defaultOptions: {
        queries: {
          retry: false,
          staleTime: Infinity,
          enabled: false
        },
      },
    });

    this.queryCache = new QueryCacheMock(this.queryKey, this.queryState);
    
  }

  getQueryCache(): QueryCache {
    return this.queryCache;
  }  
}

const mockQuery = <T>(queryKey: QueryKey, data: T): QueryClient => {
  return new QueryClientMock(queryKey, {
    status: "success",
    data,
  } as QueryState<T>);
};

const mockQueryLoading = <T>(queryKey: QueryKey): QueryClient => {
  return new QueryClientMock(queryKey, { status: "pending" } as QueryState<T>);
};

const mockQueryError = <T>(queryKey: QueryKey, error: Error): QueryClient => {
  return new QueryClientMock(queryKey, {
    status: "error",
    error,
    fetchFailureCount: 1,
    fetchStatus: "idle",
  } as QueryState<T>);
};


export { mockQuery, mockQueryLoading, mockQueryError };

useGetData.ts

import { useQuery } from "@tanstack/react-query";

export function useGetData() {
    return useQuery({["queryKey"], queryFn: async () => {
      // your implementation to fetch data
    }});
}