[Refactor] UI - Policies: Extract usePolicyVersions hook from FlowBuilderPage

Replace useEffect + useState fetch pattern for policy version management
with React Query hooks (useQuery + useMutation), following established
codebase conventions.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
yuneng-jiang 2026-03-20 16:25:40 -07:00
parent c4aa15b4e2
commit 79aea5ddcf
3 changed files with 461 additions and 76 deletions

View file

@ -0,0 +1,303 @@
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
import { renderHook, waitFor } from "@testing-library/react";
import React, { ReactNode } from "react";
import { beforeEach, describe, expect, it, vi } from "vitest";
import {
usePolicyVersions,
useCreatePolicyVersion,
useUpdatePolicyVersionStatus,
} from "./usePolicyVersions";
// ── Mocks ───────────────────────────────────────────────────────────────────
const mockListPolicyVersions = vi.fn();
const mockCreatePolicyVersion = vi.fn();
const mockUpdatePolicyVersionStatus = vi.fn();
vi.mock("@/components/networking", () => ({
listPolicyVersions: (...args: unknown[]) => mockListPolicyVersions(...args),
createPolicyVersion: (...args: unknown[]) => mockCreatePolicyVersion(...args),
updatePolicyVersionStatus: (...args: unknown[]) =>
mockUpdatePolicyVersionStatus(...args),
}));
vi.mock("@/components/molecules/notifications_manager", () => ({
default: {
success: vi.fn(),
fromBackend: vi.fn(),
},
}));
// Import the mocked module to assert on it
import NotificationsManager from "@/components/molecules/notifications_manager";
const mockUseAuthorized = vi.fn();
vi.mock("../useAuthorized", () => ({
default: () => mockUseAuthorized(),
}));
// ── Setup ───────────────────────────────────────────────────────────────────
describe("usePolicyVersions", () => {
let queryClient: QueryClient;
beforeEach(() => {
queryClient = new QueryClient({
defaultOptions: {
queries: { retry: false },
mutations: { retry: false },
},
});
vi.clearAllMocks();
mockUseAuthorized.mockReturnValue({
accessToken: "test-access-token",
userRole: "Admin",
userId: "test-user-id",
token: "test-token",
userEmail: "test@example.com",
premiumUser: false,
});
});
const wrapper = ({ children }: { children: ReactNode }) =>
React.createElement(QueryClientProvider, { client: queryClient }, children);
// ── Query tests ─────────────────────────────────────────────────────────
it("fetches versions when policyName is provided", async () => {
const mockResponse = {
policy_name: "my-policy",
versions: [
{ policy_id: "v1", policy_name: "my-policy", version_number: 1, version_status: "production" },
{ policy_id: "v2", policy_name: "my-policy", version_number: 2, version_status: "draft" },
],
total_count: 2,
};
mockListPolicyVersions.mockResolvedValue(mockResponse);
const { result } = renderHook(
() => usePolicyVersions({ policyName: "my-policy" }),
{ wrapper }
);
await waitFor(() => expect(result.current.isSuccess).toBe(true));
expect(mockListPolicyVersions).toHaveBeenCalledWith("test-access-token", "my-policy");
expect(result.current.data?.versions).toHaveLength(2);
expect(result.current.data?.versions[0].policy_id).toBe("v1");
});
it("does not fetch when policyName is null", () => {
const { result } = renderHook(
() => usePolicyVersions({ policyName: null }),
{ wrapper }
);
expect(result.current.fetchStatus).toBe("idle");
expect(mockListPolicyVersions).not.toHaveBeenCalled();
});
it("does not fetch when enabled is false", () => {
const { result } = renderHook(
() => usePolicyVersions({ policyName: "my-policy", enabled: false }),
{ wrapper }
);
expect(result.current.fetchStatus).toBe("idle");
expect(mockListPolicyVersions).not.toHaveBeenCalled();
});
it("defaults versions to empty array when response has undefined versions", async () => {
mockListPolicyVersions.mockResolvedValue({
policy_name: "my-policy",
versions: undefined,
total_count: 0,
});
const { result } = renderHook(
() => usePolicyVersions({ policyName: "my-policy" }),
{ wrapper }
);
await waitFor(() => expect(result.current.isSuccess).toBe(true));
expect(result.current.data?.versions).toEqual([]);
});
});
describe("useCreatePolicyVersion", () => {
let queryClient: QueryClient;
beforeEach(() => {
queryClient = new QueryClient({
defaultOptions: {
queries: { retry: false },
mutations: { retry: false },
},
});
vi.clearAllMocks();
mockUseAuthorized.mockReturnValue({
accessToken: "test-access-token",
userRole: "Admin",
userId: "test-user-id",
token: "test-token",
userEmail: "test@example.com",
premiumUser: false,
});
});
const wrapper = ({ children }: { children: ReactNode }) =>
React.createElement(QueryClientProvider, { client: queryClient }, children);
it("calls createPolicyVersion and shows success notification", async () => {
const newPolicy = { policy_id: "v3", policy_name: "my-policy", version_number: 3 };
mockCreatePolicyVersion.mockResolvedValue(newPolicy);
const { result } = renderHook(
() => useCreatePolicyVersion("my-policy"),
{ wrapper }
);
const returned = await result.current.mutateAsync();
expect(mockCreatePolicyVersion).toHaveBeenCalledWith("test-access-token", "my-policy");
expect(returned).toEqual(newPolicy);
expect(NotificationsManager.success).toHaveBeenCalledWith("New draft version created");
});
it("shows error notification on failure", async () => {
mockCreatePolicyVersion.mockRejectedValue(new Error("Server error"));
const { result } = renderHook(
() => useCreatePolicyVersion("my-policy"),
{ wrapper }
);
await expect(result.current.mutateAsync()).rejects.toThrow("Server error");
expect(NotificationsManager.fromBackend).toHaveBeenCalledWith(
"Failed to create version: Server error"
);
});
it("throws when policyName is null", async () => {
const { result } = renderHook(
() => useCreatePolicyVersion(null),
{ wrapper }
);
await expect(result.current.mutateAsync()).rejects.toThrow(
"Missing access token or policy name"
);
});
});
describe("useUpdatePolicyVersionStatus", () => {
let queryClient: QueryClient;
beforeEach(() => {
queryClient = new QueryClient({
defaultOptions: {
queries: { retry: false },
mutations: { retry: false },
},
});
vi.clearAllMocks();
mockUseAuthorized.mockReturnValue({
accessToken: "test-access-token",
userRole: "Admin",
userId: "test-user-id",
token: "test-token",
userEmail: "test@example.com",
premiumUser: false,
});
});
const wrapper = ({ children }: { children: ReactNode }) =>
React.createElement(QueryClientProvider, { client: queryClient }, children);
it("publishes a version and shows success notification", async () => {
const updatedPolicy = { policy_id: "v2", version_status: "published" };
mockUpdatePolicyVersionStatus.mockResolvedValue(updatedPolicy);
const { result } = renderHook(
() => useUpdatePolicyVersionStatus("my-policy"),
{ wrapper }
);
const returned = await result.current.mutateAsync({
policyId: "v2",
status: "published",
});
expect(mockUpdatePolicyVersionStatus).toHaveBeenCalledWith(
"test-access-token",
"v2",
"published"
);
expect(returned).toEqual(updatedPolicy);
expect(NotificationsManager.success).toHaveBeenCalledWith(
expect.stringContaining("Version published")
);
});
it("promotes to production and shows success notification", async () => {
const updatedPolicy = { policy_id: "v2", version_status: "production" };
mockUpdatePolicyVersionStatus.mockResolvedValue(updatedPolicy);
const { result } = renderHook(
() => useUpdatePolicyVersionStatus("my-policy"),
{ wrapper }
);
await result.current.mutateAsync({
policyId: "v2",
status: "production",
});
expect(mockUpdatePolicyVersionStatus).toHaveBeenCalledWith(
"test-access-token",
"v2",
"production"
);
expect(NotificationsManager.success).toHaveBeenCalledWith(
"Version promoted to production"
);
});
it("shows error notification on publish failure", async () => {
mockUpdatePolicyVersionStatus.mockRejectedValue(new Error("Forbidden"));
const { result } = renderHook(
() => useUpdatePolicyVersionStatus("my-policy"),
{ wrapper }
);
await expect(
result.current.mutateAsync({ policyId: "v2", status: "published" })
).rejects.toThrow("Forbidden");
expect(NotificationsManager.fromBackend).toHaveBeenCalledWith(
"Failed to publish: Forbidden"
);
});
it("shows error notification on promote failure", async () => {
mockUpdatePolicyVersionStatus.mockRejectedValue(new Error("Not found"));
const { result } = renderHook(
() => useUpdatePolicyVersionStatus("my-policy"),
{ wrapper }
);
await expect(
result.current.mutateAsync({ policyId: "v2", status: "production" })
).rejects.toThrow("Not found");
expect(NotificationsManager.fromBackend).toHaveBeenCalledWith(
"Failed to promote to production: Not found"
);
});
});

View file

@ -0,0 +1,123 @@
import { useQuery, useMutation, useQueryClient } from "@tanstack/react-query";
import useAuthorized from "../useAuthorized";
import { createQueryKeys } from "../common/queryKeysFactory";
import {
listPolicyVersions,
createPolicyVersion,
updatePolicyVersionStatus,
} from "@/components/networking";
import NotificationsManager from "@/components/molecules/notifications_manager";
import { Policy } from "@/components/policies/types";
// ── Query keys ──────────────────────────────────────────────────────────────
export const policyVersionKeys = createQueryKeys("policyVersions");
// ── Types ───────────────────────────────────────────────────────────────────
export interface PolicyVersionsResponse {
policy_name: string;
versions: Policy[];
total_count: number;
}
// ── Fetch function ──────────────────────────────────────────────────────────
const fetchPolicyVersions = async (
accessToken: string,
policyName: string
): Promise<PolicyVersionsResponse> => {
return await listPolicyVersions(accessToken, policyName);
};
// ── Hook ────────────────────────────────────────────────────────────────────
export interface UsePolicyVersionsOptions {
policyName: string | null | undefined;
enabled?: boolean;
}
export const usePolicyVersions = ({
policyName,
enabled = true,
}: UsePolicyVersionsOptions) => {
const { accessToken } = useAuthorized();
return useQuery<PolicyVersionsResponse>({
queryKey: policyVersionKeys.detail(policyName ?? ""),
queryFn: async () => await fetchPolicyVersions(accessToken!, policyName!),
enabled: Boolean(accessToken && policyName && enabled),
select: (data) => ({
...data,
versions: data.versions ?? [],
}),
});
};
// ── Mutations ───────────────────────────────────────────────────────────────
export const useCreatePolicyVersion = (policyName: string | null | undefined) => {
const { accessToken } = useAuthorized();
const queryClient = useQueryClient();
return useMutation<Policy, Error>({
mutationFn: async () => {
if (!accessToken || !policyName) {
throw new Error("Missing access token or policy name");
}
return await createPolicyVersion(accessToken, policyName);
},
onSuccess: () => {
NotificationsManager.success("New draft version created");
if (policyName) {
queryClient.invalidateQueries({
queryKey: policyVersionKeys.detail(policyName),
});
}
},
onError: (error) => {
NotificationsManager.fromBackend(
"Failed to create version: " + error.message
);
},
});
};
export const useUpdatePolicyVersionStatus = (
policyName: string | null | undefined
) => {
const { accessToken } = useAuthorized();
const queryClient = useQueryClient();
return useMutation<
Policy,
Error,
{ policyId: string; status: "published" | "production" }
>({
mutationFn: async ({ policyId, status }) => {
if (!accessToken) {
throw new Error("Missing access token");
}
return await updatePolicyVersionStatus(accessToken, policyId, status);
},
onSuccess: (_data, variables) => {
const label =
variables.status === "published"
? "Version published. You can test it in the Playground by selecting this version in the Policies dropdown."
: "Version promoted to production";
NotificationsManager.success(label);
if (policyName) {
queryClient.invalidateQueries({
queryKey: policyVersionKeys.detail(policyName),
});
}
},
onError: (error, variables) => {
const action =
variables.status === "published" ? "publish" : "promote to production";
NotificationsManager.fromBackend(
`Failed to ${action}: ${error.message}`
);
},
});
};

View file

@ -5,8 +5,13 @@ import { ArrowLeftIcon, PlusIcon } from "@heroicons/react/outline";
import { DotsVerticalIcon } from "@heroicons/react/solid";
import { GuardrailPipeline, PipelineStep, PipelineTestResult, PolicyCreateRequest, PolicyUpdateRequest, Policy } from "./types";
import { Guardrail } from "../guardrails/types";
import { testPipelineCall, listPolicyVersions, createPolicyVersion, updatePolicyVersionStatus } from "../networking";
import { testPipelineCall } from "../networking";
import NotificationsManager from "../molecules/notifications_manager";
import {
usePolicyVersions,
useCreatePolicyVersion,
useUpdatePolicyVersionStatus,
} from "@/app/(dashboard)/hooks/policies/usePolicyVersions";
import {
getComplianceDatasetPrompts,
getFrameworks,
@ -1288,10 +1293,6 @@ export const FlowBuilderPage: React.FC<FlowBuilderPageProps> = ({
const [pipeline, setPipeline] = useState<GuardrailPipeline>(
() => derivePipelineFromPolicy(editingPolicy)
);
const [versions, setVersions] = useState<Policy[]>([]);
const [isVersionsLoading, setIsVersionsLoading] = useState(false);
const [isCreatingVersion, setIsCreatingVersion] = useState(false);
const [isUpdatingStatus, setIsUpdatingStatus] = useState(false);
// Sync local state when editingPolicy changes (e.g. user switched version)
React.useEffect(() => {
@ -1300,45 +1301,23 @@ export const FlowBuilderPage: React.FC<FlowBuilderPageProps> = ({
setPipeline(derivePipelineFromPolicy(editingPolicy));
}, [editingPolicy?.policy_id, editingPolicy?.policy_name, editingPolicy?.description, editingPolicy?.pipeline, editingPolicy?.guardrails_add]);
// Fetch versions when editing an existing policy by name
React.useEffect(() => {
if (!showVersionsSidebar || !editingPolicy?.policy_name || !accessToken) {
setVersions([]);
return;
}
let cancelled = false;
setIsVersionsLoading(true);
listPolicyVersions(accessToken, editingPolicy.policy_name)
.then((res) => {
if (!cancelled) setVersions(res.versions || []);
})
.catch(() => {
if (!cancelled) setVersions([]);
})
.finally(() => {
if (!cancelled) setIsVersionsLoading(false);
});
return () => {
cancelled = true;
};
}, [showVersionsSidebar, editingPolicy?.policy_name, accessToken]);
// ── Version management via React Query hooks ──────────────────────────────
const {
data: versionsData,
isPending: isVersionsLoading,
} = usePolicyVersions({
policyName: editingPolicy?.policy_name,
enabled: showVersionsSidebar,
});
const versions = versionsData?.versions ?? [];
const createVersionMutation = useCreatePolicyVersion(editingPolicy?.policy_name);
const updateStatusMutation = useUpdatePolicyVersionStatus(editingPolicy?.policy_name);
const handleNewVersion = async () => {
if (!accessToken || !editingPolicy?.policy_name) return;
setIsCreatingVersion(true);
try {
const newPolicy = await createPolicyVersion(accessToken, editingPolicy.policy_name);
NotificationsManager.success("New draft version created");
onVersionCreated?.(newPolicy);
const list = await listPolicyVersions(accessToken, editingPolicy.policy_name);
setVersions(list.versions ?? []);
} catch (error) {
NotificationsManager.fromBackend(
"Failed to create version: " + (error instanceof Error ? error.message : String(error))
);
} finally {
setIsCreatingVersion(false);
}
const newPolicy = await createVersionMutation.mutateAsync();
onVersionCreated?.(newPolicy);
};
const handleSelectVersion = (policy: Policy) => {
@ -1346,41 +1325,21 @@ export const FlowBuilderPage: React.FC<FlowBuilderPageProps> = ({
};
const handlePublishVersion = async () => {
if (!accessToken || !editingPolicy?.policy_id) return;
setIsUpdatingStatus(true);
try {
const updated = await updatePolicyVersionStatus(accessToken, editingPolicy.policy_id, "published");
NotificationsManager.success(
"Version published. You can test it in the Playground by selecting this version in the Policies dropdown."
);
const list = await listPolicyVersions(accessToken, editingPolicy.policy_name ?? "");
setVersions(list.versions ?? []);
onVersionStatusUpdated?.(updated);
} catch (error) {
NotificationsManager.fromBackend(
"Failed to publish: " + (error instanceof Error ? error.message : String(error))
);
} finally {
setIsUpdatingStatus(false);
}
if (!editingPolicy?.policy_id) return;
const updated = await updateStatusMutation.mutateAsync({
policyId: editingPolicy.policy_id,
status: "published",
});
onVersionStatusUpdated?.(updated);
};
const handlePromoteToProduction = async () => {
if (!accessToken || !editingPolicy?.policy_id) return;
setIsUpdatingStatus(true);
try {
const updated = await updatePolicyVersionStatus(accessToken, editingPolicy.policy_id, "production");
NotificationsManager.success("Version promoted to production");
const list = await listPolicyVersions(accessToken, editingPolicy.policy_name ?? "");
setVersions(list.versions ?? []);
onVersionStatusUpdated?.(updated);
} catch (error) {
NotificationsManager.fromBackend(
"Failed to promote to production: " + (error instanceof Error ? error.message : String(error))
);
} finally {
setIsUpdatingStatus(false);
}
if (!editingPolicy?.policy_id) return;
const updated = await updateStatusMutation.mutateAsync({
policyId: editingPolicy.policy_id,
status: "production",
});
onVersionStatusUpdated?.(updated);
};
const handleSave = async () => {
@ -1540,8 +1499,8 @@ export const FlowBuilderPage: React.FC<FlowBuilderPageProps> = ({
accessToken={accessToken}
versions={versions}
isLoading={isVersionsLoading}
isCreatingVersion={isCreatingVersion}
isUpdatingStatus={isUpdatingStatus}
isCreatingVersion={createVersionMutation.isPending}
isUpdatingStatus={updateStatusMutation.isPending}
onNewVersion={handleNewVersion}
onSelectVersion={handleSelectVersion}
onPublish={handlePublishVersion}