atom-with-async.test.ts

// @vitest-environment jsdom

import { act, renderHook } from "@testing-library/react";
import { useAtom } from "jotai";
import { expect, test } from "vitest";

import { atomWithAysnc } from "./atom-with-async";

const delay = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms));

const getTestAtom = () =>
  atomWithAysnc({
    defaultValue: "DEFAULT",
    fetchInitial: () => delay(10).then(() => "FETCHED"),
    setter: () => ["OPTIMISTIC", delay(10).then(() => "REAL")],
  });

test("start with DEFAULT, switch to FETCHED when ready", async () => {
  const testAtom = getTestAtom();
  const { result } = renderHook(() => useAtom(testAtom));
  expect(result.current[0]).toBe("DEFAULT");

  await delay(15);
  expect(result.current[0]).toBe("FETCHED");
});

test("set new value", async () => {
  const testAtom = getTestAtom();
  const { result } = renderHook(() => useAtom(testAtom));

  await delay(15);
  await act(() => {
    result.current[1](23);
  });
  expect(result.current[0]).toBe("OPTIMISTIC");

  await delay(15);
  expect(result.current[0]).toBe("REAL");
});

atom-with-async.ts

import { atom } from "jotai";
import { loadable } from "jotai/utils";

export function atomWithAysnc<S, V>({
  defaultValue,
  fetchInitial,
  setter,
}: {
  defaultValue: S;
  fetchInitial: () => Promise<S>;
  setter: (value: V, state: S) => [S, Promise<S>];
}) {
  const initialAtom = loadable(atom(fetchInitial));
  const stateAtom = atom<S | null>(null);
  const finalAtom = atom(
    (get) => {
      const state = get(stateAtom);
      if (state) return state;
      const initial = get(initialAtom);
      if (initial.state === "hasData") return initial.data;
      return defaultValue;
    },
    async (get, set, value: V) => {
      const [optimistic, real] = setter(value, get(finalAtom));
      set(stateAtom, optimistic);
      set(stateAtom, await real);
    },
  );
  return finalAtom;
}