chore: add proxy provider decorator for storybook (#18023)

site/src/modules/resources/AgentRow.stories.tsx

@@ -1,8 +1,12 @@
 import type { Meta, StoryObj } from "@storybook/react";
-import { ProxyContext, getPreferredProxy } from "contexts/ProxyContext";
+import { getPreferredProxy } from "contexts/ProxyContext";
 import { chromatic } from "testHelpers/chromatic";
 import * as M from "testHelpers/entities";
-import { withDashboardProvider, withWebSocket } from "testHelpers/storybook";
+import {
+	withDashboardProvider,
+	withProxyProvider,
+	withWebSocket,
+} from "testHelpers/storybook";
 import { AgentRow } from "./AgentRow";
 
 const defaultAgentMetadata = [
@@ -94,32 +98,7 @@ const meta: Meta<typeof AgentRow> = {
 		showApps: true,
 		storybookAgentMetadata: defaultAgentMetadata,
 	},
-	decorators: [
-		(Story) => (
-			<ProxyContext.Provider
-				value={{
-					proxyLatencies: M.MockProxyLatencies,
-					proxy: getPreferredProxy([], undefined),
-					proxies: [],
-					isLoading: false,
-					isFetched: true,
-					setProxy: () => {
-						return;
-					},
-					clearProxy: () => {
-						return;
-					},
-					refetchProxyLatencies: (): Date => {
-						return new Date();
-					},
-				}}
-			>
-				<Story />
-			</ProxyContext.Provider>
-		),
-		withDashboardProvider,
-		withWebSocket,
-	],
+	decorators: [withProxyProvider(), withDashboardProvider, withWebSocket],
 	parameters: {
 		chromatic,
 		queries: [
@@ -253,31 +232,13 @@ export const Off: Story = {
 
 export const ShowingPortForward: Story = {
 	decorators: [
-		(Story) => (
-			<ProxyContext.Provider
-				value={{
-					proxyLatencies: M.MockProxyLatencies,
-					proxy: getPreferredProxy(
-						M.MockWorkspaceProxies,
-						M.MockPrimaryWorkspaceProxy,
-					),
-					proxies: M.MockWorkspaceProxies,
-					isLoading: false,
-					isFetched: true,
-					setProxy: () => {
-						return;
-					},
-					clearProxy: () => {
-						return;
-					},
-					refetchProxyLatencies: (): Date => {
-						return new Date();
-					},
-				}}
-			>
-				<Story />
-			</ProxyContext.Provider>
-		),
+		withProxyProvider({
+			proxy: getPreferredProxy(
+				M.MockWorkspaceProxies,
+				M.MockPrimaryWorkspaceProxy,
+			),
+			proxies: M.MockWorkspaceProxies,
+		}),
 	],
 };
 

site/src/modules/resources/AppLink/AppLink.stories.tsx

@@ -1,48 +1,25 @@
 import type { Meta, StoryObj } from "@storybook/react";
-import { ProxyContext, getPreferredProxy } from "contexts/ProxyContext";
+import { getPreferredProxy } from "contexts/ProxyContext";
 import {
 	MockPrimaryWorkspaceProxy,
-	MockProxyLatencies,
 	MockWorkspace,
 	MockWorkspaceAgent,
 	MockWorkspaceApp,
 	MockWorkspaceProxies,
 } from "testHelpers/entities";
-import { withGlobalSnackbar } from "testHelpers/storybook";
+import { withGlobalSnackbar, withProxyProvider } from "testHelpers/storybook";
 import { AppLink } from "./AppLink";
 
 const meta: Meta<typeof AppLink> = {
 	title: "modules/resources/AppLink",
 	component: AppLink,
 	decorators: [
-		(Story) => (
-			<ProxyContext.Provider
-				value={{
-					proxyLatencies: MockProxyLatencies,
-					proxy: {
-						...getPreferredProxy(
-							MockWorkspaceProxies,
-							MockPrimaryWorkspaceProxy,
-						),
-						preferredWildcardHostname: "*.super_proxy.tld",
-					},
-					proxies: MockWorkspaceProxies,
-					isLoading: false,
-					isFetched: true,
-					setProxy: () => {
-						return;
-					},
-					clearProxy: () => {
-						return;
-					},
-					refetchProxyLatencies: (): Date => {
-						return new Date();
-					},
-				}}
-			>
-				<Story />
-			</ProxyContext.Provider>
-		),
+		withProxyProvider({
+			proxy: {
+				...getPreferredProxy(MockWorkspaceProxies, MockPrimaryWorkspaceProxy),
+				preferredWildcardHostname: "*.super_proxy.tld",
+			},
+		}),
 	],
 };
 

site/src/modules/resources/ResourceCard.stories.tsx

@@ -1,19 +1,16 @@
 import type { Meta, StoryObj } from "@storybook/react";
-import type { WorkspaceAgent } from "api/typesGenerated";
-import { ProxyContext, getPreferredProxy } from "contexts/ProxyContext";
-import {
-	MockProxyLatencies,
-	MockWorkspaceResource,
-} from "testHelpers/entities";
+import { MockWorkspaceResource } from "testHelpers/entities";
+import { withProxyProvider } from "testHelpers/storybook";
 import { AgentRowPreview } from "./AgentRowPreview";
 import { ResourceCard } from "./ResourceCard";
 
 const meta: Meta<typeof ResourceCard> = {
 	title: "modules/resources/ResourceCard",
 	component: ResourceCard,
+	decorators: [withProxyProvider()],
 	args: {
 		resource: MockWorkspaceResource,
-		agentRow: getAgentRow,
+		agentRow: (agent) => <AgentRowPreview agent={agent} key={agent.id} />,
 	},
 };
 
@@ -67,31 +64,5 @@ export const BunchOfMetadata: Story = {
 				},
 			],
 		},
-		agentRow: getAgentRow,
 	},
 };
-
-function getAgentRow(agent: WorkspaceAgent): JSX.Element {
-	return (
-		<ProxyContext.Provider
-			value={{
-				proxyLatencies: MockProxyLatencies,
-				proxy: getPreferredProxy([], undefined),
-				proxies: [],
-				isLoading: false,
-				isFetched: true,
-				setProxy: () => {
-					return;
-				},
-				clearProxy: () => {
-					return;
-				},
-				refetchProxyLatencies: (): Date => {
-					return new Date();
-				},
-			}}
-		>
-			<AgentRowPreview agent={agent} key={agent.id} />
-		</ProxyContext.Provider>
-	);
-}

site/src/modules/resources/Resources.stories.tsx

@@ -1,20 +1,19 @@
 import type { Meta, StoryObj } from "@storybook/react";
-import type { WorkspaceAgent } from "api/typesGenerated";
-import { ProxyContext, getPreferredProxy } from "contexts/ProxyContext";
 import {
-	MockProxyLatencies,
 	MockWorkspaceResource,
 	MockWorkspaceResourceMultipleAgents,
 } from "testHelpers/entities";
+import { withProxyProvider } from "testHelpers/storybook";
 import { AgentRowPreview } from "./AgentRowPreview";
 import { Resources } from "./Resources";
 
 const meta: Meta<typeof Resources> = {
 	title: "modules/resources/Resources",
 	component: Resources,
+	decorators: [withProxyProvider()],
 	args: {
 		resources: [MockWorkspaceResource],
-		agentRow: getAgentRow,
+		agentRow: (agent) => <AgentRowPreview key={agent.id} agent={agent} />,
 	},
 };
 
@@ -163,31 +162,5 @@ export const BunchOfDevicesWithMetadata: Story = {
 				),
 			},
 		],
-		agentRow: getAgentRow,
 	},
 };
-
-function getAgentRow(agent: WorkspaceAgent): JSX.Element {
-	return (
-		<ProxyContext.Provider
-			value={{
-				proxyLatencies: MockProxyLatencies,
-				proxy: getPreferredProxy([], undefined),
-				proxies: [],
-				isLoading: false,
-				isFetched: true,
-				setProxy: () => {
-					return;
-				},
-				clearProxy: () => {
-					return;
-				},
-				refetchProxyLatencies: (): Date => {
-					return new Date();
-				},
-			}}
-		>
-			<AgentRowPreview key={agent.id} agent={agent} />
-		</ProxyContext.Provider>
-	);
-}

site/src/modules/workspaces/WorkspaceAppStatus/WorkspaceAppStatus.stories.tsx

@@ -1,38 +1,12 @@
 import type { Meta, StoryObj } from "@storybook/react";
-import { ProxyContext, getPreferredProxy } from "contexts/ProxyContext";
-import {
-	MockProxyLatencies,
-	MockWorkspaceAppStatus,
-} from "testHelpers/entities";
+import { MockWorkspaceAppStatus } from "testHelpers/entities";
+import { withProxyProvider } from "testHelpers/storybook";
 import { WorkspaceAppStatus } from "./WorkspaceAppStatus";
 
 const meta: Meta<typeof WorkspaceAppStatus> = {
 	title: "modules/workspaces/WorkspaceAppStatus",
 	component: WorkspaceAppStatus,
-	decorators: [
-		(Story) => (
-			<ProxyContext.Provider
-				value={{
-					proxyLatencies: MockProxyLatencies,
-					proxy: getPreferredProxy([], undefined),
-					proxies: [],
-					isLoading: false,
-					isFetched: true,
-					clearProxy: () => {
-						return;
-					},
-					setProxy: () => {
-						return;
-					},
-					refetchProxyLatencies: (): Date => {
-						return new Date();
-					},
-				}}
-			>
-				<Story />
-			</ProxyContext.Provider>
-		),
-	],
+	decorators: [withProxyProvider()],
 };
 
 export default meta;

site/src/pages/WorkspacePage/AppStatuses.stories.tsx

@@ -1,12 +1,11 @@
 import type { Meta, StoryObj } from "@storybook/react";
-import { ProxyContext, getPreferredProxy } from "contexts/ProxyContext";
 import {
-	MockProxyLatencies,
 	MockWorkspace,
 	MockWorkspaceAgent,
 	MockWorkspaceApp,
 	MockWorkspaceAppStatus,
 } from "testHelpers/entities";
+import { withProxyProvider } from "testHelpers/storybook";
 import { AppStatuses } from "./AppStatuses";
 
 const meta: Meta<typeof AppStatuses> = {
@@ -15,31 +14,7 @@ const meta: Meta<typeof AppStatuses> = {
 	args: {
 		referenceDate: new Date("2024-03-26T15:15:00Z"),
 	},
-	// Add decorator for ProxyContext
-	decorators: [
-		(Story) => (
-			<ProxyContext.Provider
-				value={{
-					proxyLatencies: MockProxyLatencies,
-					proxy: getPreferredProxy([], undefined),
-					proxies: [],
-					isLoading: false,
-					isFetched: true,
-					clearProxy: () => {
-						return;
-					},
-					setProxy: () => {
-						return;
-					},
-					refetchProxyLatencies: (): Date => {
-						return new Date();
-					},
-				}}
-			>
-				<Story />
-			</ProxyContext.Provider>
-		),
-	],
+	decorators: [withProxyProvider()],
 };
 
 export default meta;

site/src/pages/WorkspacePage/Workspace.stories.tsx

@@ -1,9 +1,12 @@
 import { action } from "@storybook/addon-actions";
 import type { Meta, StoryObj } from "@storybook/react";
 import type { ProvisionerJobLog } from "api/typesGenerated";
-import { ProxyContext, getPreferredProxy } from "contexts/ProxyContext";
 import * as Mocks from "testHelpers/entities";
-import { withAuthProvider, withDashboardProvider } from "testHelpers/storybook";
+import {
+	withAuthProvider,
+	withDashboardProvider,
+	withProxyProvider,
+} from "testHelpers/storybook";
 import type { WorkspacePermissions } from "../../modules/workspaces/permissions";
 import { Workspace } from "./Workspace";
 
@@ -43,32 +46,7 @@ const meta: Meta<typeof Workspace> = {
 		],
 		user: Mocks.MockUserOwner,
 	},
-	decorators: [
-		withAuthProvider,
-		withDashboardProvider,
-		(Story) => (
-			<ProxyContext.Provider
-				value={{
-					proxyLatencies: Mocks.MockProxyLatencies,
-					proxy: getPreferredProxy([], undefined),
-					proxies: [],
-					isLoading: false,
-					isFetched: true,
-					clearProxy: () => {
-						return;
-					},
-					setProxy: () => {
-						return;
-					},
-					refetchProxyLatencies: (): Date => {
-						return new Date();
-					},
-				}}
-			>
-				<Story />
-			</ProxyContext.Provider>
-		),
-	],
+	decorators: [withAuthProvider, withDashboardProvider, withProxyProvider()],
 };
 
 export default meta;

site/src/pages/WorkspacesPage/WorkspacesPageView.stories.tsx

@@ -10,7 +10,6 @@ import {
 	getDefaultFilterProps,
 } from "components/Filter/storyHelpers";
 import { DEFAULT_RECORDS_PER_PAGE } from "components/PaginationWidget/utils";
-import { ProxyContext, getPreferredProxy } from "contexts/ProxyContext";
 import dayjs from "dayjs";
 import uniqueId from "lodash/uniqueId";
 import type { ComponentProps } from "react";
@@ -18,7 +17,6 @@ import {
 	MockBuildInfo,
 	MockOrganization,
 	MockPendingProvisionerJob,
-	MockProxyLatencies,
 	MockStoppedWorkspace,
 	MockTemplate,
 	MockUserOwner,
@@ -27,7 +25,11 @@ import {
 	MockWorkspaceAppStatus,
 	mockApiError,
 } from "testHelpers/entities";
-import { withAuthProvider, withDashboardProvider } from "testHelpers/storybook";
+import {
+	withAuthProvider,
+	withDashboardProvider,
+	withProxyProvider,
+} from "testHelpers/storybook";
 import { WorkspacesPageView } from "./WorkspacesPageView";
 
 const createWorkspace = (
@@ -147,32 +149,7 @@ const meta: Meta<typeof WorkspacesPageView> = {
 		],
 		user: MockUserOwner,
 	},
-	decorators: [
-		withAuthProvider,
-		withDashboardProvider,
-		(Story) => (
-			<ProxyContext.Provider
-				value={{
-					proxyLatencies: MockProxyLatencies,
-					proxy: getPreferredProxy([], undefined),
-					proxies: [],
-					isLoading: false,
-					isFetched: true,
-					clearProxy: () => {
-						return;
-					},
-					setProxy: () => {
-						return;
-					},
-					refetchProxyLatencies: (): Date => {
-						return new Date();
-					},
-				}}
-			>
-				<Story />
-			</ProxyContext.Provider>
-		),
-	],
+	decorators: [withAuthProvider, withDashboardProvider, withProxyProvider()],
 };
 
 export default meta;

site/src/testHelpers/storybook.tsx

@@ -4,6 +4,11 @@ import { getAuthorizationKey } from "api/queries/authCheck";
 import { hasFirstUserKey, meKey } from "api/queries/users";
 import type { Entitlements } from "api/typesGenerated";
 import { GlobalSnackbar } from "components/GlobalSnackbar/GlobalSnackbar";
+import {
+	ProxyContext,
+	type ProxyContextValue,
+	getPreferredProxy,
+} from "contexts/ProxyContext";
 import { AuthProvider } from "contexts/auth/AuthProvider";
 import { DashboardContext } from "modules/dashboard/DashboardProvider";
 import { DeploymentConfigContext } from "modules/management/DeploymentConfigProvider";
@@ -17,6 +22,7 @@ import {
 	MockDeploymentConfig,
 	MockEntitlements,
 	MockOrganizationPermissions,
+	MockProxyLatencies,
 } from "./entities";
 
 export const withDashboardProvider = (
@@ -155,3 +161,30 @@ export const withOrganizationSettingsProvider = (Story: FC) => {
 		</OrganizationSettingsContext.Provider>
 	);
 };
+
+export const withProxyProvider =
+	(value?: Partial<ProxyContextValue>) => (Story: FC) => {
+		return (
+			<ProxyContext.Provider
+				value={{
+					proxyLatencies: MockProxyLatencies,
+					proxy: getPreferredProxy([], undefined),
+					proxies: [],
+					isLoading: false,
+					isFetched: true,
+					setProxy: () => {
+						return;
+					},
+					clearProxy: () => {
+						return;
+					},
+					refetchProxyLatencies: (): Date => {
+						return new Date();
+					},
+					...value,
+				}}
+			>
+				<Story />
+			</ProxyContext.Provider>
+		);
+	};