From 07087452edfd15c6f9a925ac17dd44bfd129eaff Mon Sep 17 00:00:00 2001
From: mikebender <mikebender@deephaven.io>
Date: Fri, 24 May 2024 11:15:32 -0400
Subject: [PATCH] Clean up based on review

- Create utility functions to make tests more readable
- Clean up some comments
---
 .../src/js/src/widget/WidgetHandler.test.tsx  | 26 +++++-----------
 .../ui/src/js/src/widget/WidgetHandler.tsx    |  2 +-
 .../ui/src/js/src/widget/WidgetTestUtils.ts   | 31 ++++++++++++++++---
 3 files changed, 35 insertions(+), 24 deletions(-)

diff --git a/plugins/ui/src/js/src/widget/WidgetHandler.test.tsx b/plugins/ui/src/js/src/widget/WidgetHandler.test.tsx
index a765a77a0..ad89d3e7e 100644
--- a/plugins/ui/src/js/src/widget/WidgetHandler.test.tsx
+++ b/plugins/ui/src/js/src/widget/WidgetHandler.test.tsx
@@ -5,9 +5,11 @@ import WidgetHandler, { WidgetHandlerProps } from './WidgetHandler';
 import { DocumentHandlerProps } from './DocumentHandler';
 import {
   makeDocumentUpdatedJsonRpcString,
-  makeSetStateResponse,
+  makeJsonRpcResponseString,
   makeWidget,
   makeWidgetDescriptor,
+  makeWidgetEventDocumentUpdated,
+  makeWidgetEventJsonRpcResponse,
 } from './WidgetTestUtils';
 
 const mockApi = { Widget: { EVENT_MESSAGE: 'message' } };
@@ -97,7 +99,7 @@ it('updates the document when event is received', async () => {
     // Respond to the setState call first
     listener({
       detail: {
-        getDataAsString: jest.fn(() => makeSetStateResponse(1, {})),
+        getDataAsString: jest.fn(() => makeJsonRpcResponseString(1, {})),
         exportedObjects: [],
       },
     });
@@ -203,22 +205,10 @@ it('updates the initial data only when fetch has changed', async () => {
   // Send the initial document
   await act(async () => {
     // Respond to the setState call first
-    listener({
-      detail: {
-        getDataAsString: jest.fn(() => makeSetStateResponse(1, {})),
-        exportedObjects: [],
-      },
-    });
+    listener(makeWidgetEventJsonRpcResponse(1));
 
     // Then send the initial document update
-    listener({
-      detail: {
-        getDataAsString: jest.fn(() =>
-          makeDocumentUpdatedJsonRpcString(document1)
-        ),
-        exportedObjects: [],
-      },
-    });
+    listener(makeWidgetEventDocumentUpdated(document1));
   });
 
   expect(mockDocumentHandler).toHaveBeenCalledWith(
@@ -261,7 +251,7 @@ it('updates the initial data only when fetch has changed', async () => {
   expect(fetch1).not.toHaveBeenCalled();
   expect(sendMessage).not.toHaveBeenCalled();
 
-  // Re-render with the fetch changed, it should set the state with the updated data
+  // Re-render with the widget descriptor changed, it should set the state with the updated data
   rerender(
     makeWidgetHandler({
       widget: widget2,
@@ -300,7 +290,7 @@ it('updates the initial data only when fetch has changed', async () => {
     // Respond to the setState call first
     listener({
       detail: {
-        getDataAsString: jest.fn(() => makeSetStateResponse(1, {})),
+        getDataAsString: jest.fn(() => makeJsonRpcResponseString(1, {})),
         exportedObjects: [],
       },
     });
diff --git a/plugins/ui/src/js/src/widget/WidgetHandler.tsx b/plugins/ui/src/js/src/widget/WidgetHandler.tsx
index 6127d1f36..918ba0bf1 100644
--- a/plugins/ui/src/js/src/widget/WidgetHandler.tsx
+++ b/plugins/ui/src/js/src/widget/WidgetHandler.tsx
@@ -67,7 +67,7 @@ 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]);
diff --git a/plugins/ui/src/js/src/widget/WidgetTestUtils.ts b/plugins/ui/src/js/src/widget/WidgetTestUtils.ts
index 1b98378e1..a452944ff 100644
--- a/plugins/ui/src/js/src/widget/WidgetTestUtils.ts
+++ b/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,14 +13,11 @@ export function makeDocumentUpdatedJsonRpc(
   };
 }
 
-export function makeSetStateResponse(
-  id: number,
-  state: Record<string, unknown>
-): string {
+export function makeJsonRpcResponseString(id: number, result = ''): string {
   return JSON.stringify({
     jsonrpc: '2.0',
     id,
-    result: '',
+    result,
   });
 }
 
@@ -29,6 +27,29 @@ export function makeDocumentUpdatedJsonRpcString(
   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',