feat(react/firestore): add useSetDocumentMutation (#158)

* feat(react/firestore): add useSetDocumentMutation

* refactor: move variables to mutation args

Author: HassanBahati

packages/react/src/firestore/index.ts

@@ -8,5 +8,6 @@ export { useDocumentQuery } from "./useDocumentQuery";
 export { useCollectionQuery } from "./useCollectionQuery";
 export { useGetAggregateFromServerQuery } from "./useGetAggregateFromServerQuery";
 export { useGetCountFromServerQuery } from "./useGetCountFromServerQuery";
+export { useSetDocumentMutation } from "./useSetDocumentMutation";
 export { useNamedQuery } from "./useNamedQuery";
 export { useDeleteDocumentMutation } from "./useDeleteDocumentMutation";

packages/react/src/firestore/useSetDocumentMutation.test.tsx

@@ -0,0 +1,165 @@
+import { renderHook, waitFor, act } from "@testing-library/react";
+import { doc, type DocumentReference, getDoc } from "firebase/firestore";
+import { beforeEach, describe, expect, test } from "vitest";
+import { useSetDocumentMutation } from "./useSetDocumentMutation";
+
+import {
+  expectFirestoreError,
+  firestore,
+  wipeFirestore,
+} from "~/testing-utils";
+import { queryClient, wrapper } from "../../utils";
+
+describe("useSetDocumentMutation", () => {
+  beforeEach(async () => {
+    await wipeFirestore();
+    queryClient.clear();
+  });
+
+  test("successfully sets a new document", async () => {
+    const docRef = doc(firestore, "tests", "setTest");
+    const testData = { foo: "bar", num: 42 };
+
+    const { result } = renderHook(() => useSetDocumentMutation(docRef), {
+      wrapper,
+    });
+
+    expect(result.current.isPending).toBe(false);
+    expect(result.current.isIdle).toBe(true);
+
+    await act(() => result.current.mutate(testData));
+
+    await waitFor(() => {
+      expect(result.current.isSuccess).toBe(true);
+    });
+
+    const snapshot = await getDoc(docRef);
+    expect(snapshot.exists()).toBe(true);
+    expect(snapshot.data()).toEqual(testData);
+  });
+
+  test("successfully overwrites existing document", async () => {
+    const docRef = doc(firestore, "tests", "overwriteTest");
+    const initialData = { foo: "initial", num: 1 };
+    const newData = { foo: "updated", num: 2 };
+
+    const { result } = renderHook(() => useSetDocumentMutation(docRef), {
+      wrapper,
+    });
+
+    await act(() => result.current.mutate(initialData));
+
+    await waitFor(() => {
+      expect(result.current.isSuccess).toBe(true);
+    });
+
+    let snapshot = await getDoc(docRef);
+    expect(snapshot.data()).toEqual(initialData);
+
+    const { result: result2 } = renderHook(
+      () => useSetDocumentMutation(docRef),
+      { wrapper }
+    );
+
+    await act(() => result2.current.mutate(newData));
+
+    await waitFor(() => {
+      expect(result2.current.isSuccess).toBe(true);
+    });
+
+    snapshot = await getDoc(docRef);
+    expect(snapshot.data()).toEqual(newData);
+  });
+
+  test("handles type-safe document data", async () => {
+    interface TestDoc {
+      foo: string;
+      num: number;
+    }
+
+    const docRef = doc(
+      firestore,
+      "tests",
+      "typedDoc"
+    ) as DocumentReference<TestDoc>;
+    const testData: TestDoc = { foo: "test", num: 123 };
+
+    const { result } = renderHook(() => useSetDocumentMutation(docRef), {
+      wrapper,
+    });
+
+    await act(() => result.current.mutate(testData));
+
+    await waitFor(() => {
+      expect(result.current.isSuccess).toBe(true);
+    });
+
+    const snapshot = await getDoc(docRef);
+    const data = snapshot.data();
+    expect(data?.foo).toBe("test");
+    expect(data?.num).toBe(123);
+  });
+
+  test("handles errors when setting to restricted collection", async () => {
+    const restrictedDocRef = doc(firestore, "restrictedCollection", "someDoc");
+    const testData = { foo: "bar" };
+
+    const { result } = renderHook(
+      () => useSetDocumentMutation(restrictedDocRef),
+      { wrapper }
+    );
+
+    await act(() => result.current.mutate(testData));
+
+    await waitFor(() => {
+      expect(result.current.isError).toBe(true);
+    });
+
+    expectFirestoreError(result.current.error, "permission-denied");
+  });
+
+  test("calls onSuccess callback after setting document", async () => {
+    const docRef = doc(firestore, "tests", "callbackTest");
+    const testData = { foo: "callback" };
+    let callbackCalled = false;
+
+    const { result } = renderHook(
+      () =>
+        useSetDocumentMutation(docRef, {
+          onSuccess: () => {
+            callbackCalled = true;
+          },
+        }),
+      { wrapper }
+    );
+
+    await act(() => result.current.mutate(testData));
+
+    await waitFor(() => {
+      expect(result.current.isSuccess).toBe(true);
+    });
+
+    expect(callbackCalled).toBe(true);
+    const snapshot = await getDoc(docRef);
+    expect(snapshot.data()?.foo).toBe("callback");
+  });
+
+  test("handles empty data object", async () => {
+    const docRef = doc(firestore, "tests", "emptyDoc");
+    const emptyData = {};
+
+    const { result } = renderHook(() => useSetDocumentMutation(docRef), {
+      wrapper,
+    });
+
+    await act(() => result.current.mutate(emptyData));
+
+    await waitFor(() => {
+      expect(result.current.isSuccess).toBe(true);
+    });
+
+    const snapshot = await getDoc(docRef);
+    expect(snapshot.exists()).toBe(true);
+    expect(snapshot.data()).toEqual({});
+  });
+});

packages/react/src/firestore/useSetDocumentMutation.ts

@@ -0,0 +1,30 @@
+import { useMutation, type UseMutationOptions } from "@tanstack/react-query";
+import {
+  type DocumentReference,
+  type FirestoreError,
+  type WithFieldValue,
+  type DocumentData,
+  setDoc,
+} from "firebase/firestore";
+
+type FirestoreUseMutationOptions<
+  TData = unknown,
+  TError = Error,
+  AppModelType extends DocumentData = DocumentData
+> = Omit<
+  UseMutationOptions<TData, TError, WithFieldValue<AppModelType>>,
+  "mutationFn"
+>;
+
+export function useSetDocumentMutation<
+  AppModelType extends DocumentData = DocumentData,
+  DbModelType extends DocumentData = DocumentData
+>(
+  documentRef: DocumentReference<AppModelType, DbModelType>,
+  options?: FirestoreUseMutationOptions<void, FirestoreError, AppModelType>
+) {
+  return useMutation<void, FirestoreError, WithFieldValue<AppModelType>>({
+    ...options,
+    mutationFn: (data) => setDoc(documentRef, data),
+  });
+}