From f9b552ddff52d8f9d316bfe3f2e16745ec55b5e3 Mon Sep 17 00:00:00 2001
From: Maxim Kholod <maxim.kholod@elastic.co>
Date: Tue, 5 Dec 2023 18:52:39 +0100
Subject: [PATCH] cover edit/upgrade flow

---
 .../fleet_extensions/policy_template_form.tsx | 12 +++++++++-
 .../setup_technology_selector.tsx             |  7 +++++-
 .../use_setup_technology.ts                   | 24 +++++++++++++++----
 3 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/x-pack/plugins/cloud_security_posture/public/components/fleet_extensions/policy_template_form.tsx b/x-pack/plugins/cloud_security_posture/public/components/fleet_extensions/policy_template_form.tsx
index 5c30aa0013bf7..65c9a83bd2ea0 100644
--- a/x-pack/plugins/cloud_security_posture/public/components/fleet_extensions/policy_template_form.tsx
+++ b/x-pack/plugins/cloud_security_posture/public/components/fleet_extensions/policy_template_form.tsx
@@ -21,6 +21,7 @@ import {
   EuiTitle,
 } from '@elastic/eui';
 import type { NewPackagePolicy } from '@kbn/fleet-plugin/public';
+import { SetupTechnology } from '@kbn/fleet-plugin/public';
 import { FormattedMessage } from '@kbn/i18n-react';
 import type {
   NewPackagePolicyInput,
@@ -531,7 +532,11 @@ export const CspPolicyTemplateForm = memo<PackagePolicyReplaceDefineStepExtensio
       agentPolicy,
       agentlessPolicy,
       handleSetupTechnologyChange,
+      isEditPage,
     });
+    const shouldRenderAgentlessSelector =
+      (isAgentlessAvailable && !isEditPage) ||
+      (isEditPage && setupTechnology === SetupTechnology.AGENTLESS);
 
     const updatePolicy = useCallback(
       (updatedPolicy: NewPackagePolicy) => {
@@ -590,6 +595,10 @@ export const CspPolicyTemplateForm = memo<PackagePolicyReplaceDefineStepExtensio
 
     // TODO: get rid of useEffect in favor of directly setting aws.credentials.type
     useEffect(() => {
+      if (isEditPage) {
+        return;
+      }
+
       setEnabledPolicyInput(input.type);
       // eslint-disable-next-line react-hooks/exhaustive-deps
     }, [setupTechnology]);
@@ -731,8 +740,9 @@ export const CspPolicyTemplateForm = memo<PackagePolicyReplaceDefineStepExtensio
           onChange={(field, value) => updatePolicy({ ...newPolicy, [field]: value })}
         />
 
-        {isAgentlessAvailable && (
+        {shouldRenderAgentlessSelector && (
           <SetupTechnologySelector
+            disabled={isEditPage}
             setupTechnology={setupTechnology}
             onSetupTechnologyChange={setSetupTechnology}
           />
diff --git a/x-pack/plugins/cloud_security_posture/public/components/fleet_extensions/setup_technology_selector/setup_technology_selector.tsx b/x-pack/plugins/cloud_security_posture/public/components/fleet_extensions/setup_technology_selector/setup_technology_selector.tsx
index c3b6572281da3..349e8d1a4e489 100644
--- a/x-pack/plugins/cloud_security_posture/public/components/fleet_extensions/setup_technology_selector/setup_technology_selector.tsx
+++ b/x-pack/plugins/cloud_security_posture/public/components/fleet_extensions/setup_technology_selector/setup_technology_selector.tsx
@@ -20,9 +20,11 @@ import {
 } from '@elastic/eui';
 
 export const SetupTechnologySelector = ({
+  disabled,
   setupTechnology,
   onSetupTechnologyChange,
 }: {
+  disabled: boolean;
   setupTechnology: SetupTechnology;
   onSetupTechnologyChange: (value: SetupTechnology) => void;
 }) => {
@@ -87,9 +89,11 @@ export const SetupTechnologySelector = ({
     <>
       <EuiSpacer size="l" />
       <EuiAccordion
+        isDisabled={disabled}
+        initialIsOpen={disabled}
         id={useGeneratedHtmlId({ prefix: 'setup-type' })}
         buttonContent={
-          <EuiLink>
+          <EuiLink disabled={disabled}>
             <FormattedMessage
               id="xpack.csp.fleetIntegration.setupTechnology.advancedOptionsLabel"
               defaultMessage="Advanced options"
@@ -108,6 +112,7 @@ export const SetupTechnologySelector = ({
           }
         >
           <EuiSuperSelect
+            disabled={disabled}
             options={options}
             valueOfSelected={setupTechnology}
             placeholder={
diff --git a/x-pack/plugins/cloud_security_posture/public/components/fleet_extensions/setup_technology_selector/use_setup_technology.ts b/x-pack/plugins/cloud_security_posture/public/components/fleet_extensions/setup_technology_selector/use_setup_technology.ts
index 5b211a7c878c3..4201ffde1b2a8 100644
--- a/x-pack/plugins/cloud_security_posture/public/components/fleet_extensions/setup_technology_selector/use_setup_technology.ts
+++ b/x-pack/plugins/cloud_security_posture/public/components/fleet_extensions/setup_technology_selector/use_setup_technology.ts
@@ -15,21 +15,31 @@ export const useSetupTechnology = ({
   agentPolicy,
   agentlessPolicy,
   handleSetupTechnologyChange,
+  isEditPage,
 }: {
   input: NewPackagePolicyInput;
   agentPolicy?: AgentPolicy;
   agentlessPolicy?: AgentPolicy;
   handleSetupTechnologyChange?: (value: SetupTechnology) => void;
+  isEditPage: boolean;
 }) => {
-  const [setupTechnology, setSetupTechnology] = useState<SetupTechnology>(
-    SetupTechnology.AGENT_BASED
-  );
   const isCspmAws = input.type === CLOUDBEAT_AWS;
   const isAgentlessAvailable = Boolean(isCspmAws && agentlessPolicy);
   const agentPolicyId = agentPolicy?.id;
   const agentlessPolicyId = agentlessPolicy?.id;
+  const [setupTechnology, setSetupTechnology] = useState<SetupTechnology>(() => {
+    if (isEditPage && agentPolicyId === SetupTechnology.AGENTLESS) {
+      return SetupTechnology.AGENTLESS;
+    }
+
+    return SetupTechnology.AGENT_BASED;
+  });
 
   useEffect(() => {
+    if (isEditPage) {
+      return;
+    }
+
     if (agentPolicyId && agentPolicyId !== agentlessPolicyId) {
       /*
         handle case when agent policy is coming from outside,
@@ -45,13 +55,17 @@ export const useSetupTechnology = ({
     } else {
       setSetupTechnology(SetupTechnology.AGENT_BASED);
     }
-  }, [agentPolicyId, agentlessPolicyId, isAgentlessAvailable]);
+  }, [agentPolicyId, agentlessPolicyId, isAgentlessAvailable, isEditPage]);
 
   useEffect(() => {
+    if (isEditPage) {
+      return;
+    }
+
     if (handleSetupTechnologyChange) {
       handleSetupTechnologyChange(setupTechnology);
     }
-  }, [handleSetupTechnologyChange, setupTechnology]);
+  }, [handleSetupTechnologyChange, isEditPage, setupTechnology]);
 
   return {
     isAgentlessAvailable,