Support figma ui on, off, and once

react-figma · Nov 21, 2022 · 47ee393 · 47ee393
1 parent ae89bf5
commit 47ee393
Show file tree

Hide file tree

Showing 2 changed files with 104 additions and 0 deletions.
diff --git a/src/__tests__/postmessage.test.ts b/src/__tests__/postmessage.test.ts
@@ -47,3 +47,62 @@ describe("postMessage", () => {
     });
   });
 });
+
+describe('on("message")', () => {
+  beforeAll(() => {
+    jest.useFakeTimers();
+  });
+
+  beforeEach(() => {
+    // @ts-ignore
+    global.figma = createFigma({});
+    // @ts-ignore
+    global.parent.postMessage = createParentPostMessage(global.figma);
+  });
+
+  it("can add a listener for a message", () => {
+    const cb = jest.fn();
+    figma.ui.on("message", cb);
+
+    parent.postMessage({ pluginMessage: "abc" }, "*");
+    parent.postMessage({ pluginMessage: "def" }, "*");
+
+    jest.runAllTimers();
+
+    expect(cb).toHaveBeenCalledTimes(2);
+    expect(cb).toHaveBeenCalledWith("abc", expect.any(Object));
+    expect(cb).toHaveBeenCalledWith("def", expect.any(Object));
+  });
+
+  it("can remove a listener for a message", () => {
+    const cb = jest.fn();
+    figma.ui.on("message", cb);
+
+    parent.postMessage({ pluginMessage: "abc" }, "*");
+
+    jest.runAllTimers();
+
+    expect(cb).toHaveBeenCalledTimes(1);
+    expect(cb).toHaveBeenCalledWith("abc", expect.any(Object));
+
+    cb.mockClear();
+    figma.ui.off("message", cb);
+
+    parent.postMessage({ pluginMessage: "def" }, "*");
+
+    expect(cb).not.toHaveBeenCalled();
+  });
+
+  it("can call a listener once", () => {
+    const cb = jest.fn();
+    figma.ui.once("message", cb);
+
+    parent.postMessage({ pluginMessage: "abc" }, "*");
+    parent.postMessage({ pluginMessage: "def" }, "*");
+
+    jest.runAllTimers();
+
+    expect(cb).toHaveBeenCalledTimes(1);
+    expect(cb).toHaveBeenCalledWith("abc", expect.any(Object));
+  });
+});
diff --git a/src/stubs.ts b/src/stubs.ts
@@ -144,8 +144,41 @@ export const createFigma = (paramConfig: TConfig): PluginAPI => {
   };
 
   class UIAPIStub {
+    _listeners = new Set<MessageEventHandler>();
+
     onmessage: MessageEventHandler | undefined;
 
+    on: (type: "message", cb: MessageEventHandler | undefined) => void = (
+      type,
+      cb
+    ) => {
+      if (type === "message" && cb) {
+        this._listeners.add(cb);
+      }
+    };
+
+    off: (type: "message", cb: MessageEventHandler | undefined) => void = (
+      type,
+      cb
+    ) => {
+      if (type === "message" && cb) {
+        this._listeners.delete(cb);
+      }
+    };
+
+    once: (type: "message", cb: MessageEventHandler | undefined) => void = (
+      type,
+      cb
+    ) => {
+      if (type === "message" && cb) {
+        const wrappedCb = (pluginMessage, props) => {
+          cb(pluginMessage, props);
+          this.off("message", wrappedCb);
+        };
+        this.on("message", wrappedCb);
+      }
+    };
+
     postMessage(pluginMessage: any, options?: UIPostMessageOptions): void {
       const message = {
         data: { pluginMessage, pluginId: "000000000000000000" },
@@ -528,5 +561,17 @@ export const createParentPostMessage = (
     } else {
       setTimeout(call, 0);
     }
+  } else {
+    const call = () => {
+      // @ts-ignore
+      figma.ui._listeners.forEach((cb: MessageEventHandler) => {
+        cb(message.pluginMessage, { origin: null });
+      });
+    };
+    if (isWithoutTimeout) {
+      call();
+    } else {
+      setTimeout(call, 0);
+    }
   }
 };