coder · BrunoQuaresma · Mar 7, 2024 · Mar 5, 2024 · Mar 5, 2024 · Mar 6, 2024
diff --git a/site/src/@types/storybook.d.ts b/site/src/@types/storybook.d.ts
@@ -3,12 +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?: {
-      messages: string[];
-    };
+    webSocket?: WebSocketEvent[];
   }
 }
diff --git a/site/src/contexts/auth/AuthProvider.tsx b/site/src/contexts/auth/AuthProvider.tsx
@@ -48,7 +48,6 @@ export const AuthContext = createContext<AuthContextValue | undefined>(
 export const AuthProvider: FC<PropsWithChildren> = ({ children }) => {
   const queryClient = useQueryClient();
   const meOptions = me();
-
   const userQuery = useQuery(meOptions);
   const authMethodsQuery = useQuery(authMethods());
   const hasFirstUserQuery = useQuery(hasFirstUser());

diff --git a/site/src/pages/CreateTemplatePage/BuildLogsDrawer.stories.tsx b/site/src/pages/CreateTemplatePage/BuildLogsDrawer.stories.tsx
@@ -46,8 +46,9 @@ export const Logs: Story = {
   },
   decorators: [withWebSocket],
   parameters: {
-    webSocket: {
-      messages: MockWorkspaceBuildLogs.map((log) => JSON.stringify(log)),
-    },
+    webSocket: MockWorkspaceBuildLogs.map((log) => ({
+      event: "message",
+      data: JSON.stringify(log),
+    })),
   },
 };
diff --git a/site/src/pages/TerminalPage/TerminalPage.stories.tsx b/site/src/pages/TerminalPage/TerminalPage.stories.tsx
@@ -0,0 +1,144 @@
+import type { Meta, StoryObj } from "@storybook/react";
+import {
+  reactRouterOutlet,
+  reactRouterParameters,
+} from "storybook-addon-react-router-v6";
+import { getAuthorizationKey } from "api/queries/authCheck";
+import { workspaceByOwnerAndNameKey } from "api/queries/workspaces";
+import type { Workspace, WorkspaceAgentLifecycle } from "api/typesGenerated";
+import { AuthProvider } from "contexts/auth/AuthProvider";
+import { permissionsToCheck } from "contexts/auth/permissions";
+import { RequireAuth } from "contexts/auth/RequireAuth";
+import {
+  MockAppearanceConfig,
+  MockAuthMethodsAll,
+  MockBuildInfo,
+  MockEntitlements,
+  MockExperiments,
+  MockUser,
+  MockWorkspace,
+  MockWorkspaceAgent,
+} from "testHelpers/entities";
+import { withWebSocket } from "testHelpers/storybook";
+import TerminalPage from "./TerminalPage";
+
+const createWorkspaceWithAgent = (lifecycle: WorkspaceAgentLifecycle) => {
+  return {
+    key: workspaceByOwnerAndNameKey(
+      MockWorkspace.owner_name,
+      MockWorkspace.name,
+    ),
+    data: {
+      ...MockWorkspace,
+      latest_build: {
+        ...MockWorkspace.latest_build,
+        resources: [
+          {
+            ...MockWorkspace.latest_build.resources[0],
+            agents: [{ ...MockWorkspaceAgent, lifecycle_state: lifecycle }],
+          },
+        ],
+      },
+    } satisfies Workspace,
+  };
+};
+
+const meta = {
+  title: "pages/Terminal",
+  component: RequireAuth,
+  parameters: {
+    layout: "fullscreen",
+    reactRouter: reactRouterParameters({
+      location: {
+        pathParams: {
+          username: `@${MockWorkspace.owner_name}`,
+          workspace: MockWorkspace.name,
+        },
+      },
+      routing: reactRouterOutlet(
+        {
+          path: `/:username/:workspace/terminal`,
+        },
+        <TerminalPage />,
+      ),
+    }),
+    queries: [
+      { key: ["me"], data: MockUser },
+      { key: ["authMethods"], data: MockAuthMethodsAll },
+      { key: ["hasFirstUser"], data: true },
+      { key: ["buildInfo"], data: MockBuildInfo },
+      { key: ["entitlements"], data: MockEntitlements },
+      { key: ["experiments"], data: MockExperiments },
+      { key: ["appearance"], data: MockAppearanceConfig },
+      {
+        key: getAuthorizationKey({ checks: permissionsToCheck }),
+        data: { editWorkspaceProxies: true },
+      },
+    ],
+  },
+  decorators: [
+    (Story) => (
+      <AuthProvider>
+        <Story />
+      </AuthProvider>
+    ),
+  ],
+} satisfies Meta<typeof TerminalPage>;
+
+export default meta;
+type Story = StoryObj<typeof TerminalPage>;
+
+export const Starting: Story = {
+  decorators: [withWebSocket],
+  parameters: {
+    ...meta.parameters,
+    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, createWorkspaceWithAgent("starting")],
+  },
+};
+
+export const Ready: Story = {
+  decorators: [withWebSocket],
+  parameters: {
+    ...meta.parameters,
+    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, createWorkspaceWithAgent("ready")],
+  },
+};
+
+export const StartError: Story = {
+  decorators: [withWebSocket],
+  parameters: {
+    ...meta.parameters,
+    webSocket: [],
+    queries: [
+      ...meta.parameters.queries,
+      createWorkspaceWithAgent("start_error"),
+    ],
+  },
+};
+
+export const ConnectionError: Story = {
+  decorators: [withWebSocket],
+  parameters: {
+    ...meta.parameters,
+    webSocket: [
+      {
+        event: "error",
+      },
+    ],
+    queries: [...meta.parameters.queries, createWorkspaceWithAgent("ready")],
+  },
+};
diff --git a/site/src/testHelpers/storybook.tsx b/site/src/testHelpers/storybook.tsx
@@ -45,25 +45,39 @@ export const withDashboardProvider = (
   );
 };
 
+type MessageEvent = Record<"data", string>;
+type CallbackFn = (ev?: MessageEvent) => void;
+
 export const withWebSocket = (Story: FC, { parameters }: StoryContext) => {
-  if (!parameters.webSocket) {
-    console.warn(
-      "Looks like you forgot to add websocket messages to the story",
-    );
+  const events = parameters.webSocket;
+
+  if (!events) {
+    console.warn("You forgot to add `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") {
-          parameters.webSocket?.messages.forEach((message) => {
-            callback({ data: message });
-          });
-        }
+      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: () => {},
     };