diff --git a/ui/components/providers/wizard/steps/credentials-step.test.tsx b/ui/components/providers/wizard/steps/credentials-step.test.tsx
new file mode 100644
index 0000000000..a45f8a5f3b
--- /dev/null
+++ b/ui/components/providers/wizard/steps/credentials-step.test.tsx
@@ -0,0 +1,79 @@
+import { render, screen } from "@testing-library/react";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+
+import { useProviderWizardStore } from "@/store/provider-wizard/store";
+import { PROVIDER_WIZARD_MODE } from "@/types/provider-wizard";
+
+import { CredentialsStep } from "./credentials-step";
+
+vi.mock("../../workflow/forms", () => ({
+ AddViaCredentialsForm: () =>
add-via-credentials-form
,
+ AddViaRoleForm: () => add-via-role-form
,
+ UpdateViaCredentialsForm: () => update-via-credentials-form
,
+ UpdateViaRoleForm: () => update-via-role-form
,
+}));
+
+vi.mock("../../workflow/forms/select-credentials-type/aws", () => ({
+ SelectViaAWS: () => select-via-aws
,
+}));
+
+vi.mock("../../workflow/forms/select-credentials-type/alibabacloud", () => ({
+ SelectViaAlibabaCloud: () => select-via-alibabacloud
,
+}));
+
+vi.mock("../../workflow/forms/select-credentials-type/cloudflare", () => ({
+ SelectViaCloudflare: () => select-via-cloudflare
,
+}));
+
+vi.mock("../../workflow/forms/select-credentials-type/gcp", () => ({
+ AddViaServiceAccountForm: () => add-via-service-account-form
,
+ SelectViaGCP: () => select-via-gcp
,
+}));
+
+vi.mock("../../workflow/forms/select-credentials-type/github", () => ({
+ SelectViaGitHub: () => select-via-github
,
+}));
+
+vi.mock("../../workflow/forms/select-credentials-type/m365", () => ({
+ SelectViaM365: () => select-via-m365
,
+}));
+
+vi.mock("../../workflow/forms/update-via-service-account-key-form", () => ({
+ UpdateViaServiceAccountForm: () => (
+ update-via-service-account-form
+ ),
+}));
+
+describe("CredentialsStep", () => {
+ beforeEach(() => {
+ sessionStorage.clear();
+ localStorage.clear();
+ useProviderWizardStore.getState().reset();
+ });
+
+ it("renders update role form when secret already exists in add mode", () => {
+ // Given
+ useProviderWizardStore.setState({
+ providerId: "provider-1",
+ providerType: "aws",
+ providerUid: "111111111111",
+ providerAlias: "Production",
+ via: "role",
+ secretId: "secret-1",
+ mode: PROVIDER_WIZARD_MODE.ADD,
+ });
+
+ // When
+ render(
+ ,
+ );
+
+ // Then
+ expect(screen.getByText("update-via-role-form")).toBeInTheDocument();
+ expect(screen.queryByText("add-via-role-form")).not.toBeInTheDocument();
+ });
+});
diff --git a/ui/components/providers/wizard/steps/credentials-step.tsx b/ui/components/providers/wizard/steps/credentials-step.tsx
index 0cd60c56d5..48c6418d46 100644
--- a/ui/components/providers/wizard/steps/credentials-step.tsx
+++ b/ui/components/providers/wizard/steps/credentials-step.tsx
@@ -4,7 +4,6 @@ import { useEffect, useState } from "react";
import { getProviderFormType } from "@/lib/provider-helpers";
import { useProviderWizardStore } from "@/store/provider-wizard/store";
-import { PROVIDER_WIZARD_MODE } from "@/types/provider-wizard";
import { ProviderType } from "@/types/providers";
import {
@@ -39,7 +38,7 @@ export function CredentialsStep({
onBack,
onFooterChange,
}: CredentialsStepProps) {
- const { providerId, providerType, providerUid, via, secretId, mode, setVia } =
+ const { providerId, providerType, providerUid, via, secretId, setVia } =
useProviderWizardStore();
const [isFormLoading, setIsFormLoading] = useState(false);
const [isFormValid, setIsFormValid] = useState(false);
@@ -50,8 +49,7 @@ export function CredentialsStep({
providerType && providerId
? getProviderFormType(providerType, via || undefined)
: null;
- const shouldUseUpdateForms =
- mode === PROVIDER_WIZARD_MODE.UPDATE && Boolean(secretId);
+ const shouldUseUpdateForms = Boolean(secretId);
const handleBack = () => {
if (via) {
diff --git a/ui/components/providers/wizard/steps/test-connection-step.test.tsx b/ui/components/providers/wizard/steps/test-connection-step.test.tsx
new file mode 100644
index 0000000000..cd4d2d725f
--- /dev/null
+++ b/ui/components/providers/wizard/steps/test-connection-step.test.tsx
@@ -0,0 +1,67 @@
+import { render, screen, waitFor } from "@testing-library/react";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+
+import { useProviderWizardStore } from "@/store/provider-wizard/store";
+import { PROVIDER_WIZARD_MODE } from "@/types/provider-wizard";
+
+import { TestConnectionStep } from "./test-connection-step";
+
+const { getProviderMock } = vi.hoisted(() => ({
+ getProviderMock: vi.fn(),
+}));
+
+vi.mock("@/actions/providers", () => ({
+ getProvider: getProviderMock,
+}));
+
+vi.mock("../../workflow/forms/test-connection-form", () => ({
+ TestConnectionForm: () => ,
+}));
+
+describe("TestConnectionStep", () => {
+ beforeEach(() => {
+ sessionStorage.clear();
+ localStorage.clear();
+ getProviderMock.mockReset();
+ useProviderWizardStore.getState().reset();
+ });
+
+ it("stores provider secret id after loading provider data", async () => {
+ // Given
+ useProviderWizardStore.setState({
+ providerId: "provider-1",
+ providerType: "aws",
+ mode: PROVIDER_WIZARD_MODE.ADD,
+ });
+ getProviderMock.mockResolvedValue({
+ data: {
+ id: "provider-1",
+ attributes: {
+ uid: "111111111111",
+ provider: "aws",
+ alias: "Production",
+ connection: { connected: false, last_checked_at: null },
+ scanner_args: {},
+ },
+ relationships: {
+ secret: { data: { type: "provider-secrets", id: "secret-1" } },
+ },
+ },
+ });
+
+ // When
+ render(
+ ,
+ );
+
+ // Then
+ await waitFor(() => {
+ expect(screen.getByTestId("test-connection-form")).toBeInTheDocument();
+ });
+ expect(useProviderWizardStore.getState().secretId).toBe("secret-1");
+ });
+});
diff --git a/ui/components/providers/wizard/steps/test-connection-step.tsx b/ui/components/providers/wizard/steps/test-connection-step.tsx
index 0c1ffea896..4c87e98d81 100644
--- a/ui/components/providers/wizard/steps/test-connection-step.tsx
+++ b/ui/components/providers/wizard/steps/test-connection-step.tsx
@@ -27,7 +27,8 @@ export function TestConnectionStep({
onResetCredentials,
onFooterChange,
}: TestConnectionStepProps) {
- const { providerId, providerType, mode } = useProviderWizardStore();
+ const { providerId, providerType, mode, setSecretId } =
+ useProviderWizardStore();
const [providerData, setProviderData] =
useState(null);
const [isLoadingProvider, setIsLoadingProvider] = useState(true);
@@ -62,11 +63,15 @@ export function TestConnectionStep({
setErrorMessage(
response.errors[0]?.detail || "Failed to load provider.",
);
+ setSecretId(null);
setProviderData(null);
setIsLoadingProvider(false);
return;
}
+ const resolvedSecretId =
+ response?.data?.relationships?.secret?.data?.id ?? null;
+ setSecretId(resolvedSecretId);
setProviderData(response as TestConnectionProviderData);
setIsLoadingProvider(false);
}
@@ -76,7 +81,7 @@ export function TestConnectionStep({
return () => {
isMounted = false;
};
- }, [providerId, providerType]);
+ }, [providerId, providerType, setSecretId]);
useEffect(() => {
const canSubmit = !isLoadingProvider && !errorMessage && !!providerData;