@marimo-team/islands 0.25.3-dev19 → 0.25.3-dev191

package/src/core/cells/__tests__/outputs.test.tsx

@@ -2,116 +2,134 @@
 
 import { act, renderHook } from "@testing-library/react";
 import { createStore, Provider } from "jotai";
-import { beforeEach, describe, expect, it, vi } from "vitest";
-import { MockNotebook } from "@/__mocks__/notebook";
-import { notebookAtom } from "@/core/cells/cells";
+import { afterEach, describe, expect, it } from "vitest";
+import { mockDocumentModel } from "@/__mocks__/document";
 import { CellId } from "@/core/cells/ids";
+import { DocumentProvider } from "@/core/notebook/document";
+import {
+  cell,
+  CONFIG,
+  document as documentOf,
+  withCells,
+} from "@/core/notebook/document/__tests__/fixtures";
 import { type AppMode, viewStateAtom } from "@/core/mode";
 import { useExpandedConsoleOutput, useExpandedOutput } from "../outputs";
 
-const saveCellConfig = vi.fn().mockResolvedValue(null);
+const stops: Array<() => void> = [];
 
-vi.mock("@/core/network/requests", () => ({
-  getRequestClient: () => ({ saveCellConfig }),
-}));
+afterEach(() => {
+  for (const stop of stops.splice(0)) {
+    stop();
+  }
+});
 
 function setup(opts: { mode: AppMode; expandOutput?: boolean }) {
   const cellId = CellId.create();
   const store = createStore();
-  store.set(
-    notebookAtom,
-    MockNotebook.notebookState({
-      cellData: {
-        [cellId]: { config: { expand_output: opts.expandOutput ?? false } },
+  const { model, stop } = mockDocumentModel(
+    store,
+    withCells(
+      documentOf(),
+      {
+        c1: null,
+        c2: null,
+        [cellId]: cell(cellId, "", {
+          config: { ...CONFIG, expand_output: opts.expandOutput ?? false },
+        }),
       },
-    }),
+      [cellId],
+    ),
   );
+  stops.push(stop);
   store.set(viewStateAtom, { mode: opts.mode, cellAnchor: null });
 
-  const { result } = renderHook(() => useExpandedOutput(cellId), {
-    wrapper: ({ children }) => <Provider store={store}>{children}</Provider>,
-  });
+  const wrapper = ({ children }: { children: React.ReactNode }) => (
+    <Provider store={store}>
+      <DocumentProvider model={model}>{children}</DocumentProvider>
+    </Provider>
+  );
+  const expanded = () =>
+    store.get(model.documentAtom)?.cells.get(cellId)?.config.expand_output;
 
-  return { cellId, store, result };
+  return { cellId, store, wrapper, expanded };
 }
 
-const configOf = (store: ReturnType<typeof createStore>, cellId: CellId) =>
-  store.get(notebookAtom).cellData[cellId].config;
-
 describe("useExpandedOutput", () => {
-  beforeEach(() => {
-    vi.clearAllMocks();
-  });
-
   it("reads the initial state from the cell's expand_output config", () => {
-    expect(setup({ mode: "edit" }).result.current[0]).toBe(false);
-    expect(setup({ mode: "edit", expandOutput: true }).result.current[0]).toBe(
-      true,
-    );
+    const collapsed = setup({ mode: "edit" });
+    expect(
+      renderHook(() => useExpandedOutput(collapsed.cellId), {
+        wrapper: collapsed.wrapper,
+      }).result.current[0],
+    ).toBe(false);
+
+    const expanded = setup({ mode: "edit", expandOutput: true });
+    expect(
+      renderHook(() => useExpandedOutput(expanded.cellId), {
+        wrapper: expanded.wrapper,
+      }).result.current[0],
+    ).toBe(true);
   });
 
   it("persists the toggle to the cell's config in edit mode", () => {
-    const { cellId, store, result } = setup({ mode: "edit" });
+    const { cellId, wrapper, expanded } = setup({ mode: "edit" });
+    const { result } = renderHook(() => useExpandedOutput(cellId), {
+      wrapper,
+    });
 
     act(() => result.current[1](true));
 
     expect(result.current[0]).toBe(true);
-    expect(configOf(store, cellId).expand_output).toBe(true);
-    expect(saveCellConfig).toHaveBeenCalledWith({
-      configs: { [cellId]: { expand_output: true } },
-    });
+    expect(expanded()).toBe(true);
   });
 
   it("collapses by clearing the config", () => {
-    const { cellId, store, result } = setup({
+    const { cellId, wrapper, expanded } = setup({
       mode: "edit",
       expandOutput: true,
     });
+    const { result } = renderHook(() => useExpandedOutput(cellId), {
+      wrapper,
+    });
 
     act(() => result.current[1](false));
 
     expect(result.current[0]).toBe(false);
-    expect(configOf(store, cellId).expand_output).toBe(false);
-    expect(saveCellConfig).toHaveBeenCalledWith({
-      configs: { [cellId]: { expand_output: false } },
-    });
+    expect(expanded()).toBe(false);
   });
 
   it("keeps the toggle in memory outside of edit mode", () => {
-    const { cellId, store, result } = setup({ mode: "read" });
+    const { cellId, wrapper, expanded } = setup({ mode: "read" });
+    const { result } = renderHook(() => useExpandedOutput(cellId), {
+      wrapper,
+    });
 
     act(() => result.current[1](true));
 
     expect(result.current[0]).toBe(true);
     // Cell configs are read-only outside of edit mode.
-    expect(configOf(store, cellId).expand_output).toBe(false);
-    expect(saveCellConfig).not.toHaveBeenCalled();
+    expect(expanded()).toBe(false);
   });
 
   it("still honors expand_output outside of edit mode", () => {
-    const { result } = setup({ mode: "read", expandOutput: true });
+    const { cellId, wrapper } = setup({ mode: "read", expandOutput: true });
+    const { result } = renderHook(() => useExpandedOutput(cellId), {
+      wrapper,
+    });
     expect(result.current[0]).toBe(true);
   });
 
   // `expand_output` covers the cell's output area only. Console output is
   // clamped independently and stays an in-memory, per-session toggle.
   it("does not touch the config when console output is expanded", () => {
-    const cellId = CellId.create();
-    const store = createStore();
-    store.set(
-      notebookAtom,
-      MockNotebook.notebookState({ cellData: { [cellId]: {} } }),
-    );
-    store.set(viewStateAtom, { mode: "edit", cellAnchor: null });
-
+    const { cellId, wrapper, expanded } = setup({ mode: "edit" });
     const { result } = renderHook(() => useExpandedConsoleOutput(cellId), {
-      wrapper: ({ children }) => <Provider store={store}>{children}</Provider>,
+      wrapper,
     });
 
     act(() => result.current[1](true));
 
     expect(result.current[0]).toBe(true);
-    expect(configOf(store, cellId).expand_output).toBe(false);
-    expect(saveCellConfig).not.toHaveBeenCalled();
+    expect(expanded()).toBe(false);
   });
 });

package/src/core/cells/__tests__/session.test.ts

@@ -46,10 +46,17 @@ describe("notebookStateFromSession", () => {
 
   const createNotebookCell = (
     id: string,
-    code: string | null = null,
-    name: string | null = null,
-    config: NotebookCell["config"] | null = null,
-    code_hash: string | null = null,
+    {
+      code = null,
+      name = null,
+      config = null,
+      code_hash = null,
+    }: {
+      code?: string | null;
+      name?: string | null;
+      config?: NotebookCell["config"] | null;
+      code_hash?: string | null;
+    } = {},
   ): NotebookCell => ({
     id,
     code,
@@ -87,7 +94,7 @@ describe("notebookStateFromSession", () => {
 
     it("logs error when session and notebook have different cells", () => {
       const session = createSession([createSessionCell("cell-1")]);
-      const notebook = createNotebook([createNotebookCell("cell-2")]);
+      const notebook = createNotebook([createNotebookCell("cell-2", {})]);
 
       const result = notebookStateFromSession(session, notebook);
 
@@ -328,7 +335,10 @@ describe("notebookStateFromSession", () => {
   describe("notebook only scenarios", () => {
     it("creates state from notebook with single cell", () => {
       const notebook = createNotebook([
-        createNotebookCell("cell-1", "print('hello')", "my_cell"),
+        createNotebookCell("cell-1", {
+          code: "print('hello')",
+          name: "my_cell",
+        }),
       ]);
       const result = notebookStateFromSession(null, notebook);
 
@@ -356,11 +366,15 @@ describe("notebookStateFromSession", () => {
 
     it("creates state from notebook with custom config", () => {
       const notebook = createNotebook([
-        createNotebookCell("cell-1", "code", "name", {
-          hide_code: true,
-          disabled: true,
-          column: 1,
-          expand_output: true,
+        createNotebookCell("cell-1", {
+          code: "code",
+          name: "name",
+          config: {
+            hide_code: true,
+            disabled: true,
+            column: 1,
+            expand_output: true,
+          },
         }),
       ]);
       const result = notebookStateFromSession(null, notebook);
@@ -464,7 +478,10 @@ describe("notebookStateFromSession", () => {
         ),
       ]);
       const notebook = createNotebook([
-        createNotebookCell("cell-1", "print('hello')", "my_cell"),
+        createNotebookCell("cell-1", {
+          code: "print('hello')",
+          name: "my_cell",
+        }),
       ]);
       const result = notebookStateFromSession(session, notebook);
 
@@ -495,8 +512,8 @@ describe("notebookStateFromSession", () => {
         createSessionCell("cell-2"),
       ]);
       const notebook = createNotebook([
-        createNotebookCell("cell-2"),
-        createNotebookCell("cell-1"),
+        createNotebookCell("cell-2", {}),
+        createNotebookCell("cell-1", {}),
       ]);
 
       const result = notebookStateFromSession(session, notebook);
@@ -514,8 +531,8 @@ describe("notebookStateFromSession", () => {
         createSessionCell("cell-2"),
       ]);
       const notebook = createNotebook([
-        createNotebookCell("cell-1"),
-        createNotebookCell("cell-2"),
+        createNotebookCell("cell-1", {}),
+        createNotebookCell("cell-2", {}),
       ]);
 
       const result = notebookStateFromSession(session, notebook);
@@ -669,12 +686,42 @@ describe("notebookStateFromSession", () => {
       };
       // Create notebook cells with code values 'a' through 'f'
       const notebookCells = [
-        createNotebookCell("cell-a", "a", null, null, hashes.a),
-        createNotebookCell("cell-b", "b", null, null, hashes.b),
-        createNotebookCell("cell-c", "c", null, null, hashes.c),
-        createNotebookCell("cell-d", "d", null, null, hashes.d),
-        createNotebookCell("cell-e", "e", null, null, hashes.e),
-        createNotebookCell("cell-f", "f", null, null, hashes.f),
+        createNotebookCell("cell-a", {
+          code: "a",
+          name: null,
+          config: null,
+          code_hash: hashes.a,
+        }),
+        createNotebookCell("cell-b", {
+          code: "b",
+          name: null,
+          config: null,
+          code_hash: hashes.b,
+        }),
+        createNotebookCell("cell-c", {
+          code: "c",
+          name: null,
+          config: null,
+          code_hash: hashes.c,
+        }),
+        createNotebookCell("cell-d", {
+          code: "d",
+          name: null,
+          config: null,
+          code_hash: hashes.d,
+        }),
+        createNotebookCell("cell-e", {
+          code: "e",
+          name: null,
+          config: null,
+          code_hash: hashes.e,
+        }),
+        createNotebookCell("cell-f", {
+          code: "f",
+          name: null,
+          config: null,
+          code_hash: hashes.f,
+        }),
       ];
       // Create session cells with matching code_hashes
       const sessionCells = [
@@ -787,11 +834,36 @@ describe("notebookStateFromSession", () => {
 
       // Notebook has cells with code 'a', 'c', 'z', 'e', 'g'
       const notebookCells = [
-        createNotebookCell("cell-a", "a", null, null, hashes.a),
-        createNotebookCell("cell-c", "c", null, null, hashes.c),
-        createNotebookCell("cell-z", "z", null, null, hashes.z),
-        createNotebookCell("cell-e", "e", null, null, hashes.e),
-        createNotebookCell("cell-g", "g", null, null, hashes.g),
+        createNotebookCell("cell-a", {
+          code: "a",
+          name: null,
+          config: null,
+          code_hash: hashes.a,
+        }),
+        createNotebookCell("cell-c", {
+          code: "c",
+          name: null,
+          config: null,
+          code_hash: hashes.c,
+        }),
+        createNotebookCell("cell-z", {
+          code: "z",
+          name: null,
+          config: null,
+          code_hash: hashes.z,
+        }),
+        createNotebookCell("cell-e", {
+          code: "e",
+          name: null,
+          config: null,
+          code_hash: hashes.e,
+        }),
+        createNotebookCell("cell-g", {
+          code: "g",
+          name: null,
+          config: null,
+          code_hash: hashes.g,
+        }),
       ];
 
       const session = createSession(sessionCells);
@@ -908,15 +980,24 @@ describe("notebookStateFromSession", () => {
 
       // Notebook has cells: cell-1 (1 / 0, no hash), cell-2 (mo.md("Hello"), hash=moMd), cell-3 (slider, hash=slider)
       const notebookCells = [
-        createNotebookCell("cell-1", "1 / 0", null, null, null), // no hash
-        createNotebookCell("cell-2", 'mo.md("Hello")', null, null, hashes.moMd),
-        createNotebookCell(
-          "cell-3",
-          "x = mo.ui.slider(0, 10)",
-          null,
-          null,
-          hashes.slider,
-        ),
+        createNotebookCell("cell-1", {
+          code: "1 / 0",
+          name: null,
+          config: null,
+          code_hash: null,
+        }), // no hash
+        createNotebookCell("cell-2", {
+          code: 'mo.md("Hello")',
+          name: null,
+          config: null,
+          code_hash: hashes.moMd,
+        }),
+        createNotebookCell("cell-3", {
+          code: "x = mo.ui.slider(0, 10)",
+          name: null,
+          config: null,
+          code_hash: hashes.slider,
+        }),
       ];
 
       const session = createSession(sessionCells);

package/src/core/cells/__tests__/with-document-save.test.ts

@@ -0,0 +1,67 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { atom } from "jotai";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { createDocumentModel } from "@/core/notebook/document";
+import { document } from "@/core/notebook/document/__tests__/fixtures";
+import type { Document } from "@/core/notebook/replica/resources/document";
+import {
+  type FakeEdit,
+  fakeBackend,
+} from "@/core/notebook/replica/__tests__/fixtures";
+import { activeDocumentModelAtom } from "@/core/session/active-notebook";
+import { store } from "@/core/state/jotai";
+import { Deferred } from "@/utils/Deferred";
+import { withDocumentSave } from "../document";
+
+const lifetimes: AbortController[] = [];
+
+afterEach(() => {
+  lifetimes.splice(0).forEach((controller) => controller.abort());
+  store.set(activeDocumentModelAtom, null);
+});
+
+describe("withDocumentSave", () => {
+  it("sends an edit made during the save only after the save finishes", async () => {
+    const source = atom<Document | null>(document());
+    const base = document();
+    const accepted: Document = {
+      ...base,
+      revision: { ...base.revision, version: 2, etag: "document-2" },
+    };
+    // Answer the edit, then publish the accepted document as the host would.
+    const edit = vi.fn<FakeEdit>().mockImplementation(async () => {
+      queueMicrotask(() => store.set(source, accepted));
+      return {
+        ok: true,
+        status: 200,
+        data: { document: accepted.revision, cells: {} },
+      };
+    });
+    const model = createDocumentModel({
+      store,
+      documentAtom: source,
+      backendAtom: atom(fakeBackend({ edit })),
+    });
+    const controller = new AbortController();
+    lifetimes.push(controller);
+    model.start({ signal: controller.signal });
+    store.set(activeDocumentModelAtom, model);
+
+    const saving = new Deferred<void>();
+    let flushing: Promise<unknown> | undefined;
+    const saved = withDocumentSave(async () => {
+      store.set(model.editAtom, (tx) => tx.setCode("c1", "x = 2"));
+      flushing = model.flush();
+      await Promise.resolve();
+      expect(edit).not.toHaveBeenCalled();
+      await saving.promise;
+      return "saved";
+    });
+
+    saving.resolve();
+    await expect(saved).resolves.toBe("saved");
+    await flushing;
+    expect(edit).toHaveBeenCalledTimes(1);
+  });
+});