[Cloud Security] fix fleet form save button bug (#211563)

## Summary

This PR tries to fix the following
[https://github.com/elastic/security-team/issues/11881](url)
the bug causes a lot of flakiness in our test cases.

### Checklist
- [x] [Unit or functional
tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html)
were updated or added to match the most common scenarios
- [x] [Flaky Test
Runner](https://ci-stats.kibana.dev/trigger_flaky_test_runner/1) was
used on any tests changed

### Closes
this PR closes the above mentioned issues in relation for this ticket -
https://github.com/elastic/security-team/issues/11881

### Video recording

https://github.com/user-attachments/assets/b6389216-8078-4f06-9a39-41b9559f8f1b

x-pack/platform/plugins/shared/fleet/public/applications/fleet/sections/agent_policy/create_package_policy_page/single_page_layout/index.tsx

@@ -122,6 +122,14 @@ export const CreatePackagePolicySinglePage: CreatePackagePolicyParams = ({
   );
 
   const [withSysMonitoring, setWithSysMonitoring] = useState<boolean>(true);
+  /*
+   * if there is no extension - will remain undefined
+   * if there is an extension and it is loaded - will be set to true, otherwise false
+   */
+  const [isFleetExtensionLoaded, setIsFleetExtensionLoaded] = useState<boolean | undefined>(
+    undefined
+  );
+
   const validation = agentPolicyFormValidation(newAgentPolicy, {
     allowedNamespacePrefixes: spaceSettings.allowedNamespacePrefixes,
   });
@@ -266,8 +274,9 @@ export const CreatePackagePolicySinglePage: CreatePackagePolicyParams = ({
   const handleExtensionViewOnChange = useCallback<
     PackagePolicyEditExtensionComponentProps['onChange']
   >(
-    ({ isValid, updatedPolicy }) => {
+    ({ isValid, updatedPolicy, isExtensionLoaded }) => {
       updatePackagePolicy(updatedPolicy);
+      setIsFleetExtensionLoaded(isExtensionLoaded);
       setFormState((prevState) => {
         if (prevState === 'VALID' && !isValid) {
           return 'INVALID';
@@ -652,7 +661,10 @@ export const CreatePackagePolicySinglePage: CreatePackagePolicyParams = ({
                         onClick={() => onSubmit()}
                         isLoading={formState === 'LOADING'}
                         disabled={
-                          formState !== 'VALID' || hasAgentPolicyError || !validationResults
+                          formState !== 'VALID' ||
+                          hasAgentPolicyError ||
+                          !validationResults ||
+                          isFleetExtensionLoaded === false
                         }
                         iconType="save"
                         color="primary"

x-pack/platform/plugins/shared/fleet/public/types/ui_extensions.ts

@@ -75,6 +75,7 @@ export interface PackagePolicyEditExtensionComponentProps {
     isValid: boolean;
     /** The updated Integration Policy to be merged back and included in the API call */
     updatedPolicy: Partial<NewPackagePolicy>;
+    isExtensionLoaded?: boolean;
   }) => void;
 }
 

x-pack/solutions/security/plugins/cloud_security_posture/public/components/fleet_extensions/policy_template_form.test.tsx

@@ -124,12 +124,14 @@ describe('<CspPolicyTemplateForm />', () => {
     edit = false,
     packageInfo = {} as PackageInfo,
     isAgentlessEnabled,
+    integrationToEnable,
   }: {
     edit?: boolean;
     newPolicy: NewPackagePolicy;
     packageInfo?: PackageInfo;
     onChange?: jest.Mock<void, [NewPackagePolicy]>;
     isAgentlessEnabled?: boolean;
+    integrationToEnable?: string;
   }) => {
     const { AppWrapper: FleetAppWrapper } = createFleetTestRendererMock();
     return (
@@ -143,6 +145,7 @@ describe('<CspPolicyTemplateForm />', () => {
               packageInfo={packageInfo}
               isEditPage={true}
               isAgentlessEnabled={isAgentlessEnabled}
+              integrationToEnable={integrationToEnable}
             />
           )}
           {!edit && (
@@ -152,6 +155,7 @@ describe('<CspPolicyTemplateForm />', () => {
               packageInfo={packageInfo}
               isEditPage={false}
               isAgentlessEnabled={isAgentlessEnabled}
+              integrationToEnable={integrationToEnable}
             />
           )}
         </TestProvider>
@@ -425,6 +429,89 @@ describe('<CspPolicyTemplateForm />', () => {
     });
   });
 
+  it('KSPM - calls onChange with isExtensionLoaded the second time after increment of package version', () => {
+    const policy = getMockPolicyK8s();
+
+    // enable all inputs of a policy template, same as fleet does
+    policy.inputs = policy.inputs.map((input) => ({
+      ...input,
+      enabled: input.policy_template === 'kspm',
+    }));
+    policy.name = 'cloud_security_posture-1';
+
+    (useParams as jest.Mock).mockReturnValue({
+      integration: 'kspm',
+    });
+
+    (useCspSetupStatusApi as jest.Mock).mockImplementation(() =>
+      createReactQueryResponseWithRefetch({
+        status: 'success',
+        data: {
+          kspm: { status: 'not-deployed', healthyAgents: 0, installedPackagePolicies: 1 },
+        },
+      })
+    );
+
+    (usePackagePolicyList as jest.Mock).mockImplementation(() =>
+      createReactQueryResponseWithRefetch({
+        status: 'success',
+        data: {
+          items: [getPosturePolicy(getMockPolicyEKS(), CLOUDBEAT_EKS)],
+        },
+      })
+    );
+
+    render(
+      <WrappedComponent
+        newPolicy={policy}
+        packageInfo={{ name: 'kspm' } as PackageInfo}
+        onChange={onChange}
+        integrationToEnable="kspm"
+      />
+    );
+
+    // 1st call happens on mount and selects the default policy template enabled input
+    expect(onChange).nthCalledWith(1, {
+      isExtensionLoaded: undefined,
+      isValid: true,
+      updatedPolicy: {
+        ...getMockPolicyK8s(),
+        name: 'cloud_security_posture-1',
+      },
+    });
+
+    // 2nd call happens on mount and increments kspm template enabled input
+    expect(onChange).nthCalledWith(2, {
+      isExtensionLoaded: undefined,
+      isValid: true,
+      updatedPolicy: {
+        ...getMockPolicyK8s(),
+        inputs: policy.inputs.map((input) => ({
+          ...input,
+          enabled: input.type === 'cloudbeat/cis_k8s',
+        })),
+        name: 'cloud_security_posture-1',
+      },
+    });
+
+    /*
+      3rd call happens when policies are fetched and the package version is incremented 
+      in that case isExtensionLoaded is set to 'true'
+    */
+    expect(onChange).nthCalledWith(3, {
+      isExtensionLoaded: true,
+      isValid: true,
+      updatedPolicy: {
+        ...getMockPolicyK8s(),
+        inputs: policy.inputs.map((input) => ({
+          ...input,
+          enabled: input.policy_template === 'kspm',
+        })),
+        name: 'kspm-1',
+      },
+    });
+  });
+
   it('selects default VULN_MGMT input selector', () => {
     const policy = getMockPolicyVulnMgmtAWS();
     // enable all inputs of a policy template, same as fleet does

x-pack/solutions/security/plugins/cloud_security_posture/public/components/fleet_extensions/policy_template_form.tsx

@@ -197,7 +197,7 @@ const AwsAccountTypeSelect = ({
 }: {
   input: Extract<NewPackagePolicyPostureInput, { type: 'cloudbeat/cis_aws' }>;
   newPolicy: NewPackagePolicy;
-  updatePolicy: (updatedPolicy: NewPackagePolicy) => void;
+  updatePolicy: (updatedPolicy: NewPackagePolicy, isExtensionLoaded?: boolean) => void;
   packageInfo: PackageInfo;
   disabled: boolean;
 }) => {
@@ -303,7 +303,7 @@ const GcpAccountTypeSelect = ({
 }: {
   input: Extract<NewPackagePolicyPostureInput, { type: 'cloudbeat/cis_gcp' }>;
   newPolicy: NewPackagePolicy;
-  updatePolicy: (updatedPolicy: NewPackagePolicy) => void;
+  updatePolicy: (updatedPolicy: NewPackagePolicy, isExtensionLoaded?: boolean) => void;
   packageInfo: PackageInfo;
   disabled: boolean;
 }) => {
@@ -443,7 +443,7 @@ const AzureAccountTypeSelect = ({
 }: {
   input: Extract<NewPackagePolicyPostureInput, { type: 'cloudbeat/cis_azure' }>;
   newPolicy: NewPackagePolicy;
-  updatePolicy: (updatedPolicy: NewPackagePolicy) => void;
+  updatePolicy: (updatedPolicy: NewPackagePolicy, isExtensionLoaded?: boolean) => void;
   disabled: boolean;
   packageInfo: PackageInfo;
   setupTechnology: SetupTechnology;
@@ -548,7 +548,7 @@ const useEnsureDefaultNamespace = ({
 }: {
   newPolicy: NewPackagePolicy;
   input: NewPackagePolicyPostureInput;
-  updatePolicy: (policy: NewPackagePolicy) => void;
+  updatePolicy: (policy: NewPackagePolicy, isExtensionLoaded?: boolean) => void;
 }) => {
   useEffect(() => {
     if (newPolicy.namespace === POSTURE_NAMESPACE) return;
@@ -572,7 +572,7 @@ const usePolicyTemplateInitialName = ({
   integration: CloudSecurityPolicyTemplate | undefined;
   newPolicy: NewPackagePolicy;
   packagePolicyList: PackagePolicy[] | undefined;
-  updatePolicy: (policy: NewPackagePolicy) => void;
+  updatePolicy: (policy: NewPackagePolicy, isExtensionLoaded?: boolean) => void;
   setCanFetchIntegration: (canFetch: boolean) => void;
 }) => {
   useEffect(() => {
@@ -586,14 +586,18 @@ const usePolicyTemplateInitialName = ({
 
     const currentIntegrationName = getMaxPackageName(integration, packagePolicyListByIntegration);
 
-    if (newPolicy.name === currentIntegrationName) {
-      return;
-    }
-
-    updatePolicy({
-      ...newPolicy,
-      name: currentIntegrationName,
-    });
+    /*
+     * If 'packagePolicyListByIntegration' is undefined it means policies were still not feteched - Array.isArray(undefined) = false
+     * if policie were fetched its an array - the check will return true
+     */
+    const isPoliciesLoaded = Array.isArray(packagePolicyListByIntegration);
+    updatePolicy(
+      {
+        ...newPolicy,
+        name: currentIntegrationName,
+      },
+      isPoliciesLoaded
+    );
     setCanFetchIntegration(false);
     // since this useEffect should only run on initial mount updatePolicy and newPolicy shouldn't re-trigger it
     // eslint-disable-next-line react-hooks/exhaustive-deps
@@ -629,7 +633,7 @@ const useCloudFormationTemplate = ({
 }: {
   packageInfo: PackageInfo;
   newPolicy: NewPackagePolicy;
-  updatePolicy: (policy: NewPackagePolicy) => void;
+  updatePolicy: (policy: NewPackagePolicy, isExtensionLoaded?: boolean) => void;
 }) => {
   useEffect(() => {
     const templateUrl = getVulnMgmtCloudFormationDefaultValue(packageInfo);
@@ -744,8 +748,8 @@ export const CspPolicyTemplateForm = memo<PackagePolicyReplaceDefineStepExtensio
     };
 
     const updatePolicy = useCallback(
-      (updatedPolicy: NewPackagePolicy) => {
-        onChange({ isValid, updatedPolicy });
+      (updatedPolicy: NewPackagePolicy, isExtensionLoaded?: boolean) => {
+        onChange({ isValid, updatedPolicy, isExtensionLoaded });
       },
       [onChange, isValid]
     );