Support error

BrunoQuaresma · BrunoQuaresma · commit 2a2f35a264ad · 2024-03-06T14:02:19.000Z
diff --git a/site/src/@types/storybook.d.ts b/site/src/@types/storybook.d.ts
@@ -7,8 +7,11 @@ declare module "@storybook/react" {
     features?: FeatureName[];
     experiments?: Experiments;
     queries?: { key: QueryKey; data: unknown }[];
-    webSocket?: {
-      messages: string[];
-    };
+    webSocket?:
+      | {
+          event: "message";
+          messages: string[];
+        }
+      | { event: "error" };
   }
 }
diff --git a/site/src/pages/CreateTemplatePage/BuildLogsDrawer.stories.tsx b/site/src/pages/CreateTemplatePage/BuildLogsDrawer.stories.tsx
@@ -47,6 +47,7 @@ export const Logs: Story = {
   decorators: [withWebSocket],
   parameters: {
     webSocket: {
+      event: "message",
       messages: MockWorkspaceBuildLogs.map((log) => JSON.stringify(log)),
     },
   },
diff --git a/site/src/pages/TerminalPage/TerminalPage.stories.tsx b/site/src/pages/TerminalPage/TerminalPage.stories.tsx
@@ -49,7 +49,6 @@ const meta = {
       { key: ["entitlements"], data: MockEntitlements },
       { key: ["experiments"], data: MockExperiments },
       { key: ["appearance"], data: MockAppearanceConfig },
-
       {
         key: getAuthorizationKey({ checks: permissionsToCheck }),
         data: { editWorkspaceProxies: true },
@@ -68,36 +67,44 @@ const meta = {
 export default meta;
 type Story = StoryObj<typeof TerminalPage>;
 
-export const Connected: Story = {
+const readyWorkspaceQuery = {
+  key: workspaceByOwnerAndNameKey(MockWorkspace.owner_name, MockWorkspace.name),
+  data: {
+    ...MockWorkspace,
+    latest_build: {
+      ...MockWorkspace.latest_build,
+      resources: [
+        {
+          ...MockWorkspace.latest_build.resources[0],
+          agents: [{ ...MockWorkspaceAgent, lifecycle_state: "ready" }],
+        },
+      ],
+    },
+  } satisfies Workspace,
+};
+
+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✗`,
       ],
     },
-    queries: [
-      ...meta.parameters.queries,
-      {
-        key: workspaceByOwnerAndNameKey(
-          MockWorkspace.owner_name,
-          MockWorkspace.name,
-        ),
-        data: {
-          ...MockWorkspace,
-          latest_build: {
-            ...MockWorkspace.latest_build,
-            resources: [
-              {
-                ...MockWorkspace.latest_build.resources[0],
-                agents: [{ ...MockWorkspaceAgent, lifecycle_state: "ready" }],
-              },
-            ],
-          },
-        } satisfies Workspace,
-      },
-    ],
+    queries: [...meta.parameters.queries, readyWorkspaceQuery],
+  },
+};
+
+export const OnError: Story = {
+  decorators: [withWebSocket],
+  parameters: {
+    ...meta.parameters,
+    webSocket: {
+      event: "error",
+    },
+    queries: [...meta.parameters.queries, readyWorkspaceQuery],
   },
 };
diff --git a/site/src/testHelpers/storybook.tsx b/site/src/testHelpers/storybook.tsx
@@ -46,24 +46,28 @@ export const withDashboardProvider = (
 };
 
 export const withWebSocket = (Story: FC, { parameters }: StoryContext) => {
-  if (!parameters.webSocket) {
-    console.warn(
-      "Looks like you forgot to add websocket messages to the story",
-    );
+  const webSocketConfig = parameters.webSocket;
+
+  if (!webSocketConfig) {
+    console.warn("Your forgot to add the `parameters.webSocket` to your story");
+    return <Story />;
   }
 
   // @ts-expect-error -- TS doesn't know about the global WebSocket
   window.WebSocket = function () {
     return {
       addEventListener: (
         type: string,
-        callback: (ev: Record<"data", string>) => void,
+        callback: (ev?: Record<"data", string>) => void,
       ) => {
-        if (type === "message") {
-          parameters.webSocket?.messages.forEach((message) => {
+        if (type === "message" && webSocketConfig.event === "message") {
+          webSocketConfig.messages.forEach((message) => {
             callback({ data: message });
           });
         }
+        if (type === "error" && webSocketConfig.event === "error") {
+          callback();
+        }
       },
       close: () => {},
     };
