- PR feedback

andrew-goldstein · andrew-goldstein · commit e6df9cc2849d · 2025-07-11T14:37:44.000-04:00
diff --git a/x-pack/solutions/security/plugins/security_solution/public/attack_discovery/pages/results/take_action/index.test.tsx b/x-pack/solutions/security/plugins/security_solution/public/attack_discovery/pages/results/take_action/index.test.tsx
@@ -8,16 +8,22 @@
 import { fireEvent, render, screen, waitFor } from '@testing-library/react';
 import React from 'react';
 
-import { SECURITY_FEATURE_ID } from '../../../../../common';
 import { useKibana } from '../../../../common/lib/kibana';
 import { TestProviders } from '../../../../common/mock';
 import { mockAttackDiscovery } from '../../mock/mock_attack_discovery';
 import { getMockAttackDiscoveryAlerts } from '../../mock/mock_attack_discovery_alerts';
+import { useAssistantAvailability } from '../../../../assistant/use_assistant_availability';
 import { TakeAction } from '.';
 
 const mockMutateAsyncBulk = jest.fn().mockResolvedValue({});
 const mockMutateAsyncStatus = jest.fn().mockResolvedValue({});
 
+jest.mock('../../../../assistant/use_assistant_availability', () => ({
+  useAssistantAvailability: jest.fn(),
+}));
+
+const mockUseAssistantAvailability = useAssistantAvailability as jest.Mock;
+
 jest.mock('../../../../common/lib/kibana', () => ({
   useKibana: jest.fn(),
 }));
@@ -54,8 +60,6 @@ jest.mock('../../utils/is_attack_discovery_alert', () => ({
 /** helper function to open the popover */
 const openPopover = () => fireEvent.click(screen.getAllByTestId('takeActionPopoverButton')[0]);
 
-const MOCK_SECURITY_FEATURE_ID = SECURITY_FEATURE_ID;
-
 const defaultProps = {
   attackDiscoveries: [mockAttackDiscovery],
   setSelectedAttackDiscoveries: jest.fn(),
@@ -73,9 +77,6 @@ describe('TakeAction', () => {
               show: true,
               save: true,
             },
-            [MOCK_SECURITY_FEATURE_ID]: {
-              configurations: undefined, // AI for SOC is not configured
-            },
           },
         },
         cases: {
@@ -104,6 +105,10 @@ describe('TakeAction', () => {
         },
       },
     });
+
+    mockUseAssistantAvailability.mockReturnValue({
+      hasSearchAILakeConfigurations: false, // AI for SOC is not configured
+    });
   });
 
   it('renders the Add to new case action', () => {
@@ -316,15 +321,13 @@ describe('TakeAction', () => {
       setSelectedAttackDiscoveries = jest.fn();
       (useKibana as jest.Mock).mockReturnValue({
         services: {
-          application: {
-            capabilities: {
-              assistant: { show: true, save: true },
-              [MOCK_SECURITY_FEATURE_ID]: { configurations: {} },
-            },
-          },
           cases: { helpers: { canUseCases: () => ({ createComment: true, read: true }) } },
         },
       });
+
+      mockUseAssistantAvailability.mockReturnValue({
+        hasSearchAILakeConfigurations: true, // AI for SOC IS configured
+      });
     });
 
     it('renders mark as closed action and takes action immediately (no modal)', async () => {
@@ -548,9 +551,6 @@ describe('TakeAction', () => {
                 show: true,
                 save: true,
               },
-              [MOCK_SECURITY_FEATURE_ID]: {
-                configurations: undefined,
-              },
             },
           },
         },
diff --git a/x-pack/solutions/security/plugins/security_solution/public/attack_discovery/pages/results/take_action/index.tsx b/x-pack/solutions/security/plugins/security_solution/public/attack_discovery/pages/results/take_action/index.tsx
@@ -21,10 +21,11 @@ import {
 } from '@elastic/eui';
 import React, { useCallback, useMemo, useState } from 'react';
 
+import { useAssistantAvailability } from '../../../../assistant/use_assistant_availability';
 import { useAddToNewCase } from './use_add_to_case';
 import { useAddToExistingCase } from './use_add_to_existing_case';
 import { useViewInAiAssistant } from '../attack_discovery_panel/view_in_ai_assistant/use_view_in_ai_assistant';
-import { APP_ID, SECURITY_FEATURE_ID } from '../../../../../common';
+import { APP_ID } from '../../../../../common';
 import { useKibana } from '../../../../common/lib/kibana';
 import * as i18n from './translations';
 import { UpdateAlertsModal } from './update_alerts_modal';
@@ -55,12 +56,9 @@ const TakeActionComponent: React.FC<Props> = ({
   );
 
   const {
-    services: {
-      application: { capabilities },
-      cases,
-    },
+    services: { cases },
   } = useKibana();
-  const aiForSoc = capabilities[SECURITY_FEATURE_ID].configurations;
+  const { hasSearchAILakeConfigurations } = useAssistantAvailability();
 
   const { attackDiscoveryAlertsEnabled } = useKibanaFeatureFlags();
 
@@ -165,12 +163,12 @@ const TakeActionComponent: React.FC<Props> = ({
 
       setPendingAction(workflowStatus);
 
-      if (aiForSoc) {
+      if (hasSearchAILakeConfigurations) {
         // there's no modal for AI for SOC, so we call onConfirm directly
         onConfirm({ updateAlerts: false, workflowStatus });
       }
     },
-    [aiForSoc, closePopover, onConfirm]
+    [closePopover, hasSearchAILakeConfigurations, onConfirm]
   );
 
   const onClickAddToNewCase = useCallback(async () => {
@@ -343,7 +341,7 @@ const TakeActionComponent: React.FC<Props> = ({
         <EuiContextMenuPanel size="s" items={allItems} />
       </EuiPopover>
 
-      {pendingAction != null && !aiForSoc && (
+      {pendingAction != null && !hasSearchAILakeConfigurations && (
         <UpdateAlertsModal
           alertsCount={alertIds.length}
           attackDiscoveriesCount={attackDiscoveryIds.length}
