@marimo-team/frontend 0.19.8-dev5 → 0.19.8-dev50

package/src/plugins/impl/anywidget/__tests__/AnyWidgetPlugin.test.tsx

@@ -1,65 +1,17 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import { act, render, waitFor } from "@testing-library/react";
+import { render, waitFor } from "@testing-library/react";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 import { TestUtils } from "@/__tests__/test-helpers";
-import type { UIElementId } from "@/core/cells/ids";
-import { MarimoIncomingMessageEvent } from "@/core/dom/events";
-import { getDirtyFields, visibleForTesting } from "../AnyWidgetPlugin";
-import { Model } from "../model";
+import type { HTMLElementNotDerivedFromRef } from "@/hooks/useEventListener";
+import { visibleForTesting } from "../AnyWidgetPlugin";
+import { MODEL_MANAGER, Model } from "../model";
+import type { WidgetModelId } from "../types";
 
 const { LoadedSlot } = visibleForTesting;
 
-describe("getDirtyFields", () => {
-  it("should return empty set when values are equal", () => {
-    const value = { foo: "bar", baz: 123 };
-    const initialValue = { foo: "bar", baz: 123 };
-
-    const result = getDirtyFields(value, initialValue);
-
-    expect(result.size).toBe(0);
-  });
-
-  it("should return keys of changed values", () => {
-    const value = { foo: "changed", baz: 123 };
-    const initialValue = { foo: "bar", baz: 123 };
-
-    const result = getDirtyFields(value, initialValue);
-
-    expect(result.size).toBe(1);
-    expect(result.has("foo")).toBe(true);
-  });
-
-  it("should handle multiple changed values", () => {
-    const value = { foo: "changed", baz: 456, extra: "new" };
-    const initialValue = { foo: "bar", baz: 123, extra: "old" };
-
-    const result = getDirtyFields(value, initialValue);
-
-    expect(result.size).toBe(3);
-    expect(result.has("foo")).toBe(true);
-    expect(result.has("baz")).toBe(true);
-    expect(result.has("extra")).toBe(true);
-  });
-
-  it("should handle nested objects correctly", () => {
-    const value = { foo: "bar", nested: { a: 1, b: 2 } };
-    const initialValue = { foo: "bar", nested: { a: 1, b: 3 } };
-
-    const result = getDirtyFields(value, initialValue);
-
-    expect(result.size).toBe(1);
-    expect(result.has("nested")).toBe(true);
-  });
-
-  it("should handle subset of initial fields", () => {
-    const value = { foo: "bar", baz: 123 };
-    const initialValue = { foo: "bar", baz: 123, full: "value" };
-
-    const result = getDirtyFields(value, initialValue);
-    expect(result.size).toBe(0);
-  });
-});
+// Helper to create typed model IDs for tests
+const asModelId = (id: string): WidgetModelId => id as WidgetModelId;
 
 // Mock a minimal AnyWidget implementation
 const mockWidget = {
@@ -76,23 +28,32 @@ vi.mock("../AnyWidgetPlugin", async () => {
 });
 
 describe("LoadedSlot", () => {
+  const modelId = asModelId("test-model-id");
+  let mockModel: Model<{ count: number }>;
+
   const mockProps = {
-    value: { count: 0 },
-    setValue: vi.fn(),
     widget: mockWidget,
-    functions: {
-      send_to_widget: vi.fn().mockResolvedValue(null),
-    },
     data: {
       jsUrl: "http://example.com/widget.js",
       jsHash: "abc123",
-      initialValue: { count: 0 },
     },
-    host: document.createElement("div"),
+    host: document.createElement(
+      "div",
+    ) as unknown as HTMLElementNotDerivedFromRef,
+    modelId: modelId,
   };
 
   beforeEach(() => {
     vi.clearAllMocks();
+    // Create and register a mock model before each test
+    mockModel = new Model(
+      { count: 0 },
+      {
+        sendUpdate: vi.fn().mockResolvedValue(undefined),
+        sendCustomMessage: vi.fn().mockResolvedValue(undefined),
+      },
+    );
+    MODEL_MANAGER.set(modelId, mockModel);
   });
 
   it("should render a div with ref", () => {
@@ -100,83 +61,36 @@ describe("LoadedSlot", () => {
     expect(container.querySelector("div")).not.toBeNull();
   });
 
-  it("should initialize model with merged values", () => {
-    const modelSpy = vi.spyOn(Model.prototype, "updateAndEmitDiffs");
-    render(<LoadedSlot {...mockProps} />);
-
-    expect(modelSpy).toHaveBeenCalledExactlyOnceWith({ count: 0 });
-  });
-
-  it("should update model when value prop changes", async () => {
-    const { rerender } = render(<LoadedSlot {...mockProps} />);
-    const modelSpy = vi.spyOn(Model.prototype, "updateAndEmitDiffs");
-
-    // Update the value prop
-    rerender(<LoadedSlot {...mockProps} value={{ count: 5 }} />);
-
-    // Model should be updated with the new value
-    expect(modelSpy).toHaveBeenCalledWith({ count: 5 });
-  });
-
-  it("should listen for incoming messages", async () => {
+  it("should call runAnyWidgetModule on initialization", async () => {
     render(<LoadedSlot {...mockProps} />);
 
-    // Send a mock message
-    const mockMessageEvent = MarimoIncomingMessageEvent.create({
-      detail: {
-        objectId: "test-id" as UIElementId,
-        message: {
-          method: "update",
-          state: { count: 10 },
-          buffer_paths: [],
-        },
-        buffers: [],
-      },
-      bubbles: false,
-      composed: true,
-    });
-    const updateAndEmitDiffsSpy = vi.spyOn(Model.prototype, "set");
-
-    // Dispatch the event on the host element
-    act(() => {
-      mockProps.host.dispatchEvent(mockMessageEvent);
-    });
-
+    // Wait a render
     await waitFor(() => {
-      expect(updateAndEmitDiffsSpy).toHaveBeenCalledWith("count", 10);
+      expect(mockWidget.render).toHaveBeenCalled();
     });
   });
 
-  it("should call runAnyWidgetModule on initialization", async () => {
+  it("should re-run widget when widget prop changes", async () => {
     const { rerender } = render(<LoadedSlot {...mockProps} />);
 
-    // Wait a render
+    // Wait for initial render
     await waitFor(() => {
       expect(mockWidget.render).toHaveBeenCalled();
     });
 
-    // Render without any prop changes
-    rerender(<LoadedSlot {...mockProps} />);
-    await TestUtils.nextTick();
+    // Create a new widget mock
+    const newMockWidget = {
+      initialize: vi.fn(),
+      render: vi.fn(),
+    };
 
-    // Still only called once
-    expect(mockWidget.render).toHaveBeenCalledTimes(1);
-
-    // Change the jsUrl
-    rerender(
-      <LoadedSlot
-        {...mockProps}
-        data={{
-          ...mockProps.data,
-          jsUrl: "http://example.com/widget-updated.js",
-        }}
-      />,
-    );
+    // Change the widget
+    rerender(<LoadedSlot {...mockProps} widget={newMockWidget} />);
     await TestUtils.nextTick();
 
-    // Wait a render
+    // Wait for re-render with new widget
     await waitFor(() => {
-      expect(mockWidget.render).toHaveBeenCalledTimes(2);
+      expect(newMockWidget.render).toHaveBeenCalled();
     });
   });
 });

package/src/plugins/impl/anywidget/__tests__/model.test.ts

@@ -9,35 +9,73 @@ import {
   vi,
 } from "vitest";
 import { TestUtils } from "@/__tests__/test-helpers";
-import type { Base64String } from "@/utils/json/base64";
 import {
-  type AnyWidgetMessage,
+  getMarimoInternal,
   handleWidgetMessage,
   Model,
   visibleForTesting,
 } from "../model";
+import type { WidgetModelId } from "../types";
 
 const { ModelManager } = visibleForTesting;
 
+// Helper to create typed model IDs for tests
+const asModelId = (id: string): WidgetModelId => id as WidgetModelId;
+
+// Mock the request client
+const mockSendModelValue = vi.fn().mockResolvedValue(null);
+vi.mock("@/core/network/requests", () => ({
+  getRequestClient: () => ({
+    sendModelValue: mockSendModelValue,
+  }),
+}));
+
+// Helper to create a mock MarimoComm
+function createMockComm<T>() {
+  return {
+    sendUpdate: vi.fn().mockResolvedValue(undefined),
+    sendCustomMessage: vi.fn().mockResolvedValue(undefined),
+  };
+}
+
 describe("Model", () => {
   let model: Model<{ foo: string; bar: number }>;
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  let onChange: (value: any) => void;
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  let sendToWidget: (req: {
-    content: unknown;
-    buffers: Base64String[];
-  }) => Promise<null | undefined>;
+  let mockComm: ReturnType<typeof createMockComm<{ foo: string; bar: number }>>;
 
   beforeEach(() => {
-    onChange = vi.fn();
-    sendToWidget = vi.fn().mockResolvedValue(null);
-    model = new Model(
-      { foo: "test", bar: 123 },
-      onChange,
-      sendToWidget,
-      new Set(),
-    );
+    mockComm = createMockComm();
+    mockSendModelValue.mockClear();
+    model = new Model({ foo: "test", bar: 123 }, mockComm);
+  });
+
+  describe("public API", () => {
+    it("should only expose AFM-compliant interface", () => {
+      // Get all enumerable own properties
+      const ownProperties = Object.keys(model).sort();
+      // Get prototype methods (excluding constructor)
+      const prototypeMethods = Object.getOwnPropertyNames(
+        Object.getPrototypeOf(model),
+      )
+        .filter((name) => name !== "constructor")
+        .sort();
+
+      // Snapshot the public API to catch accidental leaks of internal methods
+      expect({ ownProperties, prototypeMethods }).toMatchInlineSnapshot(`
+        {
+          "ownProperties": [
+            "widget_manager",
+          ],
+          "prototypeMethods": [
+            "get",
+            "off",
+            "on",
+            "save_changes",
+            "send",
+            "set",
+          ],
+        }
+      `);
+    });
   });
 
   describe("get/set", () => {
@@ -70,7 +108,7 @@ describe("Model", () => {
       model.set("bar", 456);
       model.save_changes();
 
-      expect(onChange).toHaveBeenCalledWith({
+      expect(mockComm.sendUpdate).toHaveBeenCalledWith({
         foo: "new value",
         bar: 456,
       });
@@ -80,7 +118,7 @@ describe("Model", () => {
       model.set("foo", "new value");
       model.save_changes();
 
-      expect(onChange).toHaveBeenCalledWith({
+      expect(mockComm.sendUpdate).toHaveBeenCalledWith({
         foo: "new value",
       });
 
@@ -88,17 +126,17 @@ describe("Model", () => {
       model.save_changes();
 
       // After clearing, only the newly changed field is sent
-      expect(onChange).toHaveBeenCalledWith({
+      expect(mockComm.sendUpdate).toHaveBeenCalledWith({
         bar: 456,
       });
     });
 
-    it("should not call onChange when no dirty fields", () => {
+    it("should not call sendUpdate when no dirty fields", () => {
       model.set("foo", "new value");
       model.save_changes();
-      model.save_changes(); // Second save should not call onChange
+      model.save_changes(); // Second save should not call sendUpdate
 
-      expect(onChange).toHaveBeenCalledTimes(1);
+      expect(mockComm.sendUpdate).toHaveBeenCalledTimes(1);
     });
   });
 
@@ -145,23 +183,28 @@ describe("Model", () => {
   describe("send", () => {
     it("should send message and handle callbacks", async () => {
       const callback = vi.fn();
-      model.send({ test: true }, callback);
+      await model.send({ test: true }, callback);
 
-      expect(sendToWidget).toHaveBeenCalledWith({
-        content: {
-          state: { test: true },
-          bufferPaths: [],
-        },
-        buffers: [],
-      });
-      await TestUtils.nextTick(); // flush
-      expect(callback).toHaveBeenCalledWith(null);
+      expect(mockComm.sendCustomMessage).toHaveBeenCalledWith(
+        { test: true },
+        [],
+      );
+      expect(callback).toHaveBeenCalled();
+    });
+
+    it("should convert buffers to DataViews", async () => {
+      const buffer = new ArrayBuffer(8);
+      await model.send({ test: true }, undefined, [buffer]);
+
+      expect(mockComm.sendCustomMessage).toHaveBeenCalledWith({ test: true }, [
+        expect.any(DataView),
+      ]);
     });
   });
 
   describe("widget_manager", () => {
-    const childModelId = "test-id";
-    const childModel = new Model({ foo: "test" }, vi.fn(), vi.fn(), new Set());
+    const childModelId = asModelId("test-id");
+    const childModel = new Model({ foo: "test" }, createMockComm());
     const manager = new ModelManager(10);
     let previousModelManager = Model._modelManager;
 
@@ -177,9 +220,9 @@ describe("Model", () => {
     });
 
     it("should throw error when accessing a model that is not registered", async () => {
-      await expect(model.widget_manager.get_model("random-id")).rejects.toThrow(
-        "Model not found for key: random-id",
-      );
+      await expect(
+        model.widget_manager.get_model(asModelId("random-id")),
+      ).rejects.toThrow("Model not found for key: random-id");
     });
 
     it("should return the registered model", async () => {
@@ -194,7 +237,7 @@ describe("Model", () => {
       const callback = vi.fn();
       model.on("change:foo", callback);
 
-      model.updateAndEmitDiffs({ foo: "test", bar: 456 });
+      getMarimoInternal(model).updateAndEmitDiffs({ foo: "test", bar: 456 });
       expect(callback).not.toHaveBeenCalled(); // foo didn't change
       expect(model.get("bar")).toBe(456);
     });
@@ -202,47 +245,33 @@ describe("Model", () => {
     it("should update and emit for deep changes", () => {
       const modelWithObject = new Model<{ foo: { nested: string } }>(
         { foo: { nested: "test" } },
-        onChange,
-        sendToWidget,
-        new Set(),
+        createMockComm(),
       );
       const callback = vi.fn();
       modelWithObject.on("change:foo", callback);
 
-      modelWithObject.updateAndEmitDiffs({ foo: { nested: "changed" } });
+      getMarimoInternal(modelWithObject).updateAndEmitDiffs({
+        foo: { nested: "changed" },
+      });
       expect(callback).toHaveBeenCalledTimes(1);
     });
 
     it("should emit change event for any changes", async () => {
       const callback = vi.fn();
       model.on("change", callback);
-      model.updateAndEmitDiffs({ foo: "changed", bar: 456 });
+      getMarimoInternal(model).updateAndEmitDiffs({ foo: "changed", bar: 456 });
       await TestUtils.nextTick(); // flush
       expect(callback).toHaveBeenCalledTimes(1);
     });
   });
 
-  describe("receiveCustomMessage", () => {
-    it("should handle update messages", () => {
-      model.receiveCustomMessage({
-        method: "update",
-        state: {
-          foo: "updated",
-          bar: 789,
-        },
-        buffer_paths: [],
-      });
-
-      expect(model.get("foo")).toBe("updated");
-      expect(model.get("bar")).toBe(789);
-    });
-
+  describe("emitCustomMessage", () => {
     it("should handle custom messages", () => {
       const callback = vi.fn();
       model.on("msg:custom", callback);
 
       const content = { type: "test" };
-      model.receiveCustomMessage({
+      getMarimoInternal(model).emitCustomMessage({
         method: "custom",
         content,
       });
@@ -256,7 +285,7 @@ describe("Model", () => {
 
       const content = { type: "test" };
       const buffer = new DataView(new ArrayBuffer(8));
-      model.receiveCustomMessage(
+      getMarimoInternal(model).emitCustomMessage(
         {
           method: "custom",
           content,
@@ -266,108 +295,85 @@ describe("Model", () => {
 
       expect(callback).toHaveBeenCalledWith(content, [buffer]);
     });
-
-    it("should log error for invalid messages", () => {
-      const consoleSpy = vi.spyOn(console, "error").mockImplementation(() => {
-        // noop
-      });
-      model.receiveCustomMessage({ invalid: "message" });
-
-      expect(consoleSpy).toHaveBeenCalledTimes(2);
-    });
   });
 });
 
 describe("ModelManager", () => {
   let modelManager = new ModelManager(50);
-  const handle = ({
-    modelId,
-    message,
-    buffers,
-  }: {
-    modelId: string;
-    message: AnyWidgetMessage;
-    buffers: readonly DataView[];
-  }) => {
-    return handleWidgetMessage({
-      modelId,
-      msg: message,
-      buffers,
-      modelManager,
-    });
-  };
+  const testId = asModelId("test-id");
 
   beforeEach(() => {
     // Clear the model manager before each test
     modelManager = new ModelManager(50);
+    mockSendModelValue.mockClear();
   });
 
   it("should set and get models", async () => {
-    const model = new Model({ count: 0 }, vi.fn(), vi.fn(), new Set());
-    modelManager.set("test-id", model);
-    const retrievedModel = await modelManager.get("test-id");
+    const model = new Model({ count: 0 }, createMockComm());
+    modelManager.set(testId, model);
+    const retrievedModel = await modelManager.get(testId);
     expect(retrievedModel).toBe(model);
   });
 
   it("should handle model not found", async () => {
-    await expect(modelManager.get("non-existent")).rejects.toThrow(
+    await expect(modelManager.get(asModelId("non-existent"))).rejects.toThrow(
       "Model not found for key: non-existent",
     );
   });
 
   it("should delete models", async () => {
-    const model = new Model({ count: 0 }, vi.fn(), vi.fn(), new Set());
-    modelManager.set("test-id", model);
-    modelManager.delete("test-id");
-    await expect(modelManager.get("test-id")).rejects.toThrow();
+    const model = new Model({ count: 0 }, createMockComm());
+    modelManager.set(testId, model);
+    modelManager.delete(testId);
+    await expect(modelManager.get(testId)).rejects.toThrow();
   });
 
   it("should handle widget messages", async () => {
-    const openMessage: AnyWidgetMessage = {
-      method: "open",
-      state: { count: 0 },
-      buffer_paths: [],
-    };
-
-    await handle({ modelId: "test-id", message: openMessage, buffers: [] });
-    const model = await modelManager.get("test-id");
+    await handleWidgetMessage(modelManager, {
+      model_id: testId,
+      message: {
+        method: "open",
+        state: { count: 0 },
+        buffer_paths: [],
+        buffers: [],
+      },
+    });
+    const model = await modelManager.get(testId);
     expect(model.get("count")).toBe(0);
 
-    const updateMessage: AnyWidgetMessage = {
-      method: "update",
-      state: {
-        count: 1,
+    await handleWidgetMessage(modelManager, {
+      model_id: testId,
+      message: {
+        method: "update",
+        state: { count: 1 },
+        buffer_paths: [],
+        buffers: [],
       },
-      buffer_paths: [],
-    };
-
-    await handle({ modelId: "test-id", message: updateMessage, buffers: [] });
+    });
     expect(model.get("count")).toBe(1);
   });
 
   it("should handle custom messages", async () => {
-    const model = new Model({ count: 0 }, vi.fn(), vi.fn(), new Set());
+    const model = new Model({ count: 0 }, createMockComm());
     const callback = vi.fn();
     model.on("msg:custom", callback);
-    modelManager.set("test-id", model);
+    modelManager.set(testId, model);
 
-    await handle({
-      modelId: "test-id",
-      message: { method: "custom", content: { count: 1 } },
-      buffers: [],
+    await handleWidgetMessage(modelManager, {
+      model_id: testId,
+      message: { method: "custom", content: { count: 1 }, buffers: [] },
     });
     expect(callback).toHaveBeenCalledWith({ count: 1 }, []);
   });
 
   it("should handle close messages", async () => {
-    const model = new Model({ count: 0 }, vi.fn(), vi.fn(), new Set());
-    modelManager.set("test-id", model);
+    const model = new Model({ count: 0 }, createMockComm());
+    modelManager.set(testId, model);
 
-    await handle({
-      modelId: "test-id",
+    await handleWidgetMessage(modelManager, {
+      model_id: testId,
       message: { method: "close" },
-      buffers: [],
     });
-    await expect(modelManager.get("test-id")).rejects.toThrow();
+    await expect(modelManager.get(testId)).rejects.toThrow();
   });
 });