fix: Reset state when new instance of widget created

plugins/ui/src/js/src/widget/WidgetHandler.test.tsx

@@ -5,8 +5,11 @@ import WidgetHandler, { WidgetHandlerProps } from './WidgetHandler';
 import { DocumentHandlerProps } from './DocumentHandler';
 import {
   makeDocumentUpdatedJsonRpcString,
+  makeJsonRpcResponseString,
   makeWidget,
   makeWidgetDescriptor,
+  makeWidgetEventDocumentUpdated,
+  makeWidgetEventJsonRpcResponse,
 } from './WidgetTestUtils';
 
 const mockApi = { Widget: { EVENT_MESSAGE: 'message' } };
@@ -26,8 +29,16 @@ function makeWidgetHandler({
   fetch = () => Promise.resolve(makeWidget()),
   widget = makeWidgetDescriptor(),
   onClose = jest.fn(),
+  initialData = undefined,
 }: Partial<WidgetHandlerProps> = {}) {
-  return <WidgetHandler fetch={fetch} widget={widget} onClose={onClose} />;
+  return (
+    <WidgetHandler
+      fetch={fetch}
+      widget={widget}
+      onClose={onClose}
+      initialData={initialData}
+    />
+  );
 }
 
 beforeEach(() => {
@@ -48,38 +59,77 @@ it('updates the document when event is received', async () => {
   const widget = makeWidgetDescriptor();
   const cleanup = jest.fn();
   const mockAddEventListener = jest.fn(() => cleanup);
+  const mockSendMessage = jest.fn();
+  const initialData = { state: { fiz: 'baz' } };
   const initialDocument = { foo: 'bar' };
   const widgetObject = makeWidget({
     addEventListener: mockAddEventListener,
-    getDataAsString: jest.fn(() =>
-      makeDocumentUpdatedJsonRpcString(initialDocument)
-    ),
+    getDataAsString: jest.fn(() => ''),
+    sendMessage: mockSendMessage,
   });
 
-  const { unmount } = render(makeWidgetHandler({ widget, fetch }));
+  const { unmount } = render(makeWidgetHandler({ widget, fetch, initialData }));
   expect(fetch).toHaveBeenCalledTimes(1);
   expect(mockAddEventListener).not.toHaveBeenCalled();
   expect(mockDocumentHandler).not.toHaveBeenCalled();
+  expect(mockSendMessage).not.toHaveBeenCalled();
   await act(async () => {
     fetchResolve!(widgetObject);
     await fetchPromise;
   });
 
   expect(mockAddEventListener).toHaveBeenCalledTimes(1);
+  expect(mockDocumentHandler).not.toHaveBeenCalled();
+
+  expect(mockSendMessage).toHaveBeenCalledWith(
+    JSON.stringify({
+      jsonrpc: '2.0',
+      id: 1,
+      method: 'setState',
+      params: [initialData.state],
+    }),
+    []
+  );
+
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  const listener = (mockAddEventListener.mock.calls[0] as any)[1];
+
+  // Send the initial document
+  await act(async () => {
+    // Respond to the setState call first
+    listener({
+      detail: {
+        getDataAsString: jest.fn(() => makeJsonRpcResponseString(1, {})),
+        exportedObjects: [],
+      },
+    });
+
+    // Then send the initial document update
+    listener({
+      detail: {
+        getDataAsString: jest.fn(() =>
+          makeDocumentUpdatedJsonRpcString(initialDocument)
+        ),
+        exportedObjects: [],
+      },
+    });
+  });
+
   expect(mockDocumentHandler).toHaveBeenCalledWith(
     expect.objectContaining({
       widget,
       children: initialDocument,
+      initialData,
     })
   );
 
-  mockDocumentHandler.mockClear();
+  const updatedDocument = { FOO: 'BAR' };
 
-  const updatedDocument = { fiz: 'baz' };
+  mockDocumentHandler.mockClear();
 
-  act(() => {
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    (mockAddEventListener.mock.calls[0] as any)[1]({
+  // Send the updated document
+  await act(async () => {
+    listener({
       detail: {
         getDataAsString: jest.fn(() =>
           makeDocumentUpdatedJsonRpcString(updatedDocument)
@@ -99,3 +149,172 @@ it('updates the document when event is received', async () => {
   unmount();
   expect(cleanup).toHaveBeenCalledTimes(1);
 });
+
+it('updates the initial data only when fetch has changed', async () => {
+  let fetchResolve1: (value: dh.Widget | PromiseLike<dh.Widget>) => void;
+  const fetchPromise1 = new Promise<dh.Widget>(resolve => {
+    fetchResolve1 = resolve;
+  });
+  const fetch1 = jest.fn(() => fetchPromise1);
+  const widget1 = makeWidgetDescriptor();
+  const cleanup = jest.fn();
+  const addEventListener = jest.fn(() => cleanup);
+  const sendMessage = jest.fn();
+  const onClose = jest.fn();
+  const data1 = { state: { fiz: 'baz' } };
+  const document1 = { foo: 'bar' };
+  const widgetObject1 = makeWidget({
+    addEventListener,
+    getDataAsString: jest.fn(() => ''),
+    sendMessage,
+  });
+
+  const { rerender, unmount } = render(
+    makeWidgetHandler({
+      widget: widget1,
+      fetch: fetch1,
+      initialData: data1,
+      onClose,
+    })
+  );
+  expect(fetch1).toHaveBeenCalledTimes(1);
+  expect(addEventListener).not.toHaveBeenCalled();
+  expect(mockDocumentHandler).not.toHaveBeenCalled();
+  expect(sendMessage).not.toHaveBeenCalled();
+  await act(async () => {
+    fetchResolve1!(widgetObject1);
+    await fetchPromise1;
+  });
+
+  expect(addEventListener).toHaveBeenCalledTimes(1);
+  expect(mockDocumentHandler).not.toHaveBeenCalled();
+
+  expect(sendMessage).toHaveBeenCalledWith(
+    JSON.stringify({
+      jsonrpc: '2.0',
+      id: 1,
+      method: 'setState',
+      params: [data1.state],
+    }),
+    []
+  );
+
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  let listener = (addEventListener.mock.calls[0] as any)[1];
+
+  // Send the initial document
+  await act(async () => {
+    // Respond to the setState call first
+    listener(makeWidgetEventJsonRpcResponse(1));
+
+    // Then send the initial document update
+    listener(makeWidgetEventDocumentUpdated(document1));
+  });
+
+  expect(mockDocumentHandler).toHaveBeenCalledWith(
+    expect.objectContaining({
+      widget: widget1,
+      children: document1,
+      initialData: data1,
+    })
+  );
+
+  let fetchResolve2: (value: dh.Widget | PromiseLike<dh.Widget>) => void;
+  const fetchPromise2 = new Promise<dh.Widget>(resolve => {
+    fetchResolve2 = resolve;
+  });
+  const widget2 = makeWidgetDescriptor();
+  const document2 = { FOO: 'BAR' };
+  const data2 = { state: { FIZ: 'BAZ' } };
+  const fetch2 = jest.fn(() => fetchPromise2);
+  const widgetObject2 = makeWidget({
+    addEventListener,
+    getDataAsString: jest.fn(() => ''),
+    sendMessage,
+  });
+
+  addEventListener.mockClear();
+  mockDocumentHandler.mockClear();
+  sendMessage.mockClear();
+  fetch1.mockClear();
+
+  // Re-render with just initial data change. It should not set the state again
+  rerender(
+    makeWidgetHandler({
+      widget: widget1,
+      fetch: fetch1,
+      initialData: data2,
+      onClose,
+    })
+  );
+
+  expect(fetch1).not.toHaveBeenCalled();
+  expect(sendMessage).not.toHaveBeenCalled();
+
+  // Re-render with the widget descriptor changed, it should set the state with the updated data
+  rerender(
+    makeWidgetHandler({
+      widget: widget2,
+      fetch: fetch2,
+      initialData: data2,
+      onClose,
+    })
+  );
+
+  await act(async () => {
+    fetchResolve2!(widgetObject2);
+    await fetchPromise2;
+  });
+
+  expect(fetch2).toHaveBeenCalledTimes(1);
+  // Should have been called when the widget was updated
+  expect(cleanup).toHaveBeenCalledTimes(1);
+  cleanup.mockClear();
+
+  // eslint-disable-next-line prefer-destructuring, @typescript-eslint/no-explicit-any
+  listener = (addEventListener.mock.calls[0] as any)[1];
+
+  expect(sendMessage).toHaveBeenCalledWith(
+    JSON.stringify({
+      jsonrpc: '2.0',
+      id: 1,
+      method: 'setState',
+      params: [data2.state],
+    }),
+    []
+  );
+  expect(sendMessage).toHaveBeenCalledTimes(1);
+
+  // Send the initial document
+  await act(async () => {
+    // Respond to the setState call first
+    listener({
+      detail: {
+        getDataAsString: jest.fn(() => makeJsonRpcResponseString(1, {})),
+        exportedObjects: [],
+      },
+    });
+
+    // Then send the initial document update
+    listener({
+      detail: {
+        getDataAsString: jest.fn(() =>
+          makeDocumentUpdatedJsonRpcString(document2)
+        ),
+        exportedObjects: [],
+      },
+    });
+  });
+
+  expect(mockDocumentHandler).toHaveBeenCalledWith(
+    expect.objectContaining({
+      widget: widget1,
+      children: document2,
+      initialData: data2,
+    })
+  );
+
+  expect(cleanup).not.toHaveBeenCalled();
+  unmount();
+  expect(cleanup).toHaveBeenCalledTimes(1);
+});

plugins/ui/src/js/src/widget/WidgetHandler.tsx

@@ -67,7 +67,10 @@ function WidgetHandler({
   const [widget, setWidget] = useState<dh.Widget>();
   const [document, setDocument] = useState<ReactNode>();
   const [error, setError] = useState<WidgetError>();
-  const [initialData] = useState(initialDataProp);
+
+  // We want to update the initial data if the widget changes, as we'll need to re-fetch the widget and want to start with a fresh state.
+  // eslint-disable-next-line react-hooks/exhaustive-deps
+  const initialData = useMemo(() => initialDataProp, [widget]);
 
   // When we fetch a widget, the client is then responsible for the exported objects.
   // These objects could stay alive even after the widget is closed if we wanted to,

plugins/ui/src/js/src/widget/WidgetTestUtils.ts

@@ -1,6 +1,7 @@
 import { WidgetDescriptor } from '@deephaven/dashboard';
 import { TestUtils } from '@deephaven/utils';
 import type { dh } from '@deephaven/jsapi-types';
+import { WidgetMessageEvent } from './WidgetTypes';
 
 export function makeDocumentUpdatedJsonRpc(
   document: Record<string, unknown> = {}
@@ -12,12 +13,43 @@ export function makeDocumentUpdatedJsonRpc(
   };
 }
 
+export function makeJsonRpcResponseString(id: number, result = ''): string {
+  return JSON.stringify({
+    jsonrpc: '2.0',
+    id,
+    result,
+  });
+}
+
 export function makeDocumentUpdatedJsonRpcString(
   document: Record<string, unknown> = {}
 ): string {
   return JSON.stringify(makeDocumentUpdatedJsonRpc(document));
 }
 
+export function makeWidgetEvent(data = ''): WidgetMessageEvent {
+  return new CustomEvent('message', {
+    detail: {
+      getDataAsBase64: () => '',
+      getDataAsString: () => data,
+      exportedObjects: [],
+    },
+  });
+}
+
+export function makeWidgetEventJsonRpcResponse(
+  id: number,
+  response = ''
+): WidgetMessageEvent {
+  return makeWidgetEvent(makeJsonRpcResponseString(id, response));
+}
+
+export function makeWidgetEventDocumentUpdated(
+  document: Record<string, unknown> = {}
+): WidgetMessageEvent {
+  return makeWidgetEvent(makeDocumentUpdatedJsonRpcString(document));
+}
+
 export function makeWidgetDescriptor({
   id = 'widget-id',
   type = 'widget-type',
@@ -34,10 +66,12 @@ export function makeWidget({
   addEventListener = jest.fn(() => jest.fn()),
   getDataAsString = () => makeDocumentUpdatedJsonRpcString(),
   exportedObjects = [],
+  sendMessage = jest.fn(),
 }: Partial<dh.Widget> = {}): dh.Widget {
   return TestUtils.createMockProxy<dh.Widget>({
     addEventListener,
     getDataAsString,
     exportedObjects,
+    sendMessage,
   });
 }