Stack multiple web socket events

BrunoQuaresma · BrunoQuaresma · commit 116b8d76b4d1 · 2024-03-06T14:24:46.000Z
diff --git a/site/src/@types/storybook.d.ts b/site/src/@types/storybook.d.ts
@@ -3,15 +3,11 @@ import type { QueryKey } from "react-query";
 import type { Experiments, FeatureName } from "api/typesGenerated";
 
 declare module "@storybook/react" {
+  type WebSocketEvent = { event: "message"; data: string } | { event: "error" };
   interface Parameters {
     features?: FeatureName[];
     experiments?: Experiments;
     queries?: { key: QueryKey; data: unknown }[];
-    webSocket?:
-      | {
-          event: "message";
-          messages: string[];
-        }
-      | { event: "error" };
+    webSocket?: WebSocketEvent[];
   }
 }
diff --git a/site/src/pages/CreateTemplatePage/BuildLogsDrawer.stories.tsx b/site/src/pages/CreateTemplatePage/BuildLogsDrawer.stories.tsx
@@ -46,9 +46,9 @@ export const Logs: Story = {
   },
   decorators: [withWebSocket],
   parameters: {
-    webSocket: {
+    webSocket: MockWorkspaceBuildLogs.map((log) => ({
       event: "message",
-      messages: MockWorkspaceBuildLogs.map((log) => JSON.stringify(log)),
-    },
+      data: JSON.stringify(log),
+    })),
   },
 };
diff --git a/site/src/pages/TerminalPage/TerminalPage.stories.tsx b/site/src/pages/TerminalPage/TerminalPage.stories.tsx
@@ -87,13 +87,13 @@ export const OnMessage: Story = {
   decorators: [withWebSocket],
   parameters: {
     ...meta.parameters,
-    webSocket: {
-      event: "message",
-      // Copied and pasted this from browser
-      messages: [
-        `[H[2J[1m[32m➜  [36mcoder[C[34mgit:([31mbq/refactor-web-term-notifications[34m) [33m✗`,
-      ],
-    },
+    webSocket: [
+      {
+        event: "message",
+        // Copied and pasted this from browser
+        data: `[H[2J[1m[32m➜  [36mcoder[C[34mgit:([31mbq/refactor-web-term-notifications[34m) [33m✗`,
+      },
+    ],
     queries: [...meta.parameters.queries, readyWorkspaceQuery],
   },
 };
@@ -102,9 +102,11 @@ export const OnError: Story = {
   decorators: [withWebSocket],
   parameters: {
     ...meta.parameters,
-    webSocket: {
-      event: "error",
-    },
+    webSocket: [
+      {
+        event: "error",
+      },
+    ],
     queries: [...meta.parameters.queries, readyWorkspaceQuery],
   },
 };
diff --git a/site/src/testHelpers/storybook.tsx b/site/src/testHelpers/storybook.tsx
@@ -45,29 +45,39 @@ export const withDashboardProvider = (
   );
 };
 
+type MessageEvent = Record<"data", string>;
+type CallbackFn = (ev?: MessageEvent) => void;
+
 export const withWebSocket = (Story: FC, { parameters }: StoryContext) => {
-  const webSocketConfig = parameters.webSocket;
+  const events = parameters.webSocket;
 
-  if (!webSocketConfig) {
+  if (!events) {
     console.warn("Your forgot to add the `parameters.webSocket` to your story");
     return <Story />;
   }
 
+  const listeners = new Map<string, CallbackFn>();
+  let callEventsDelay: number;
+
   // @ts-expect-error -- TS doesn't know about the global WebSocket
   window.WebSocket = function () {
     return {
-      addEventListener: (
-        type: string,
-        callback: (ev?: Record<"data", string>) => void,
-      ) => {
-        if (type === "message" && webSocketConfig.event === "message") {
-          webSocketConfig.messages.forEach((message) => {
-            callback({ data: message });
-          });
-        }
-        if (type === "error" && webSocketConfig.event === "error") {
-          callback();
-        }
+      addEventListener: (type: string, callback: CallbackFn) => {
+        listeners.set(type, callback);
+
+        // Runs when the last event listener is added
+        clearTimeout(callEventsDelay);
+        callEventsDelay = window.setTimeout(() => {
+          for (const entry of events) {
+            const callback = listeners.get(entry.event);
+
+            if (callback) {
+              entry.event === "message"
+                ? callback({ data: entry.data })
+                : callback();
+            }
+          }
+        }, 0);
       },
       close: () => {},
     };
