[UI v2] Use cva() on StateBadge (#16209)

PrefectHQ · Dec 4, 2024 · 1cab155 · 1cab155
1 parent 730f56e
commit 1cab155
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 68 deletions.
diff --git a/ui-v2/src/components/ui/state-badge/index.tsx b/ui-v2/src/components/ui/state-badge/index.tsx
@@ -1,14 +1,15 @@
-import { Badge } from "../badge";
 import type { components } from "@/api/prefect";
+import { cva } from "class-variance-authority";
 import {
+	BanIcon,
+	CheckIcon,
 	ClockIcon,
 	PauseIcon,
-	XIcon,
-	CheckIcon,
-	ServerCrashIcon,
-	BanIcon,
 	PlayIcon,
+	ServerCrashIcon,
+	XIcon,
 } from "lucide-react";
+import { Badge } from "../badge";
 
 const ICONS = {
 	COMPLETED: CheckIcon,
@@ -25,29 +26,31 @@ const ICONS = {
 	React.ElementType
 >;
 
-const CLASSES = {
-	COMPLETED: "bg-green-50 text-green-600 hover:bg-green-50",
-	FAILED: "bg-red-50 text-red-600 hover:bg-red-50",
-	RUNNING: "bg-blue-100 text-blue-700 hover:bg-blue-100",
-	CANCELLED: "bg-gray-300 text-gray-800 hover:bg-gray-300",
-	CANCELLING: "bg-gray-300 text-gray-800 hover:bg-gray-300",
-	CRASHED: "bg-orange-50 text-orange-600 hover:bg-orange-50",
-	PAUSED: "bg-gray-300 text-gray-800 hover:bg-gray-300",
-	PENDING: "bg-gray-300 text-gray-800 hover:bg-gray-300",
-	SCHEDULED: "bg-yellow-100 text-yellow-700 hover:bg-yellow-100",
-} as const satisfies Record<components["schemas"]["StateType"], string>;
+const stateBadgeVariants = cva("gap-1", {
+	variants: {
+		state: {
+			COMPLETED: "bg-green-50 text-green-600 hover:bg-green-50",
+			FAILED: "bg-red-50 text-red-600 hover:bg-red-50",
+			RUNNING: "bg-blue-100 text-blue-700 hover:bg-blue-100",
+			CANCELLED: "bg-gray-300 text-gray-800 hover:bg-gray-300",
+			CANCELLING: "bg-gray-300 text-gray-800 hover:bg-gray-300",
+			CRASHED: "bg-orange-50 text-orange-600 hover:bg-orange-50",
+			PAUSED: "bg-gray-300 text-gray-800 hover:bg-gray-300",
+			PENDING: "bg-gray-300 text-gray-800 hover:bg-gray-300",
+			SCHEDULED: "bg-yellow-100 text-yellow-700 hover:bg-yellow-100",
+		} satisfies Record<components["schemas"]["StateType"], string>,
+	},
+});
 
 export const StateBadge = ({
 	state,
 }: { state: components["schemas"]["State"] }) => {
 	const Icon = ICONS[state.type];
 	return (
-		<Badge className={CLASSES[state.type]}>
-			<div className="flex items-center gap-1">
-				<Icon size={16} />
+		<Badge className={stateBadgeVariants({ state: state.type })}>
+			<Icon size={16} />
 
-				{state.name}
-			</div>
+			{state.name}
 		</Badge>
 	);
 };
diff --git a/ui-v2/src/components/ui/state-badge/state-badge.stories.tsx b/ui-v2/src/components/ui/state-badge/state-badge.stories.tsx
@@ -1,47 +1,35 @@
+import type { components } from "@/api/prefect.ts";
 import type { Meta, StoryObj } from "@storybook/react";
 import { StateBadge } from ".";
 
-const meta = {
-	title: "UI/StateBadge",
-	component: StateBadge,
-	argTypes: {
-		state: {
-			options: [
-				"COMPLETED",
-				"FAILED",
-				"RUNNING",
-				"PENDING",
-				"PAUSED",
-				"CANCELLED",
-				"CANCELLING",
-				"CRASHED",
-				"SCHEDULED",
-				"LATE",
-			],
-			mapping: {
-				COMPLETED: { type: "COMPLETED", name: "Completed" },
-				FAILED: { type: "FAILED", name: "Failed" },
-				RUNNING: { type: "RUNNING", name: "Running" },
-				PENDING: { type: "PENDING", name: "Pending" },
-				PAUSED: { type: "PAUSED", name: "Paused" },
-				CANCELLED: { type: "CANCELLED", name: "Cancelled" },
-				CANCELLING: { type: "CANCELLING", name: "Cancelling" },
-				CRASHED: { type: "CRASHED", name: "Crashed" },
-				SCHEDULED: { type: "SCHEDULED", name: "Scheduled" },
-				LATE: { type: "SCHEDULED", name: "Late" },
-			},
-		},
-	},
-} satisfies Meta<typeof StateBadge>;
+const badgesByState: Record<components["schemas"]["StateType"], string[]> = {
+	COMPLETED: ["Completed"],
+	FAILED: ["Failed"],
+	RUNNING: ["Running"],
+	PENDING: ["Pending"],
+	PAUSED: ["Paused"],
+	CANCELLED: ["Cancelled"],
+	CANCELLING: ["Cancelling"],
+	CRASHED: ["Crashed"],
+	SCHEDULED: ["Scheduled", "Late"],
+};
 
-export default meta;
-type Story = StoryObj<typeof meta>;
+export const story: StoryObj = { name: "StateBadge" };
 
-export const States: Story = {
-	args: {
-		state: {
-			type: "COMPLETED",
-			name: "Completed",
-		},
+export default {
+	title: "UI/StateBadge",
+	component: function StateBadgeStories() {
+		return (
+			<div className="flex flex-col gap-4 items-start">
+				{Object.entries(badgesByState).map(([type, names]) =>
+					names.map((name) => (
+						<StateBadge
+							key={name}
+							state={{ type, name } as components["schemas"]["State"]}
+						/>
+					)),
+				)}
+			</div>
+		);
 	},
-};
+} satisfies Meta;
diff --git a/ui-v2/src/components/ui/state-badge/state-badge.test.tsx b/ui-v2/src/components/ui/state-badge/state-badge.test.tsx
@@ -1,15 +1,15 @@
 import { render, screen } from "@testing-library/react";
-import { StateBadge } from "./index";
 import {
+	BanIcon,
+	CheckIcon,
 	ClockIcon,
 	PauseIcon,
-	XIcon,
-	CheckIcon,
-	ServerCrashIcon,
-	BanIcon,
 	PlayIcon,
+	ServerCrashIcon,
+	XIcon,
 } from "lucide-react";
 import { describe, expect, test } from "vitest";
+import { StateBadge } from "./index";
 
 describe("StateBadge", () => {
 	const states = [
@@ -87,7 +87,7 @@ describe("StateBadge", () => {
 				SCHEDULED: "bg-yellow-100 text-yellow-700 hover:bg-yellow-100",
 			}[type];
 
-			expect(badge?.parentElement).toHaveClass(...expectedClasses.split(" "));
+			expect(badge).toHaveClass(...expectedClasses.split(" "));
 		},
 	);
 });