optimization

Signed-off-by: Qxisylolo <[email protected]>

public/components/ui_action_context_menu.test.tsx

@@ -0,0 +1,117 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+import { ActionContextMenu, Props } from './ui_action_context_menu';
+import { httpServiceMock } from '../../../../src/core/public/http/http_service.mock';
+import React from 'react';
+import { render, screen, fireEvent, waitFor } from '@testing-library/react';
+import { dataPluginMock } from '../../../../src/plugins/data/public/mocks';
+import { BehaviorSubject } from 'rxjs';
+import { AssistantClient } from '../../public/services/assistant_client';
+import { act } from 'react-dom/test-utils';
+
+jest.mock('../services', () => ({
+  getUiActions: () => ({
+    getTriggerActions: jest.fn().mockReturnValue([
+      {
+        definition: {
+          order: 1,
+          id: 'suggestAnomalyDetector',
+          type: 'suggestAnomalyDetector',
+        },
+        id: 'suggestAnomalyDetector',
+        type: 'suggestAnomalyDetector',
+        order: 1,
+      },
+    ]),
+  }),
+}));
+
+describe('ActionContextMenu', () => {
+  let mockProps: Props;
+  const resultSummaryEnabled$ = new BehaviorSubject<boolean>(true);
+  const isQuerySummaryCollapsed$ = new BehaviorSubject<boolean>(false);
+  const isSummaryAgentAvailable$ = new BehaviorSubject<boolean>(false);
+  const mockAgentConfigExists = jest.fn().mockResolvedValue({ exists: true });
+  beforeEach(() => {
+    jest.clearAllMocks();
+    mockProps = {
+      label: 'OpenSearch Assistant',
+      httpSetup: httpServiceMock.createSetupContract(),
+      data: dataPluginMock.createSetupContract(),
+      resultSummaryEnabled$,
+      isQuerySummaryCollapsed$,
+      isSummaryAgentAvailable$,
+      assistantServiceStart: {
+        client: ({
+          agentConfigExists: mockAgentConfigExists,
+          executeAgent: jest.fn(),
+          executeAgentByConfigName: jest.fn(),
+          http: {},
+        } as unknown) as AssistantClient,
+      },
+    };
+    // Mock subscribe methods to emit values synchronously
+    jest.spyOn(resultSummaryEnabled$, 'subscribe').mockImplementation((callback) => {
+      callback(resultSummaryEnabled$.value);
+      return { unsubscribe: jest.fn() };
+    });
+    jest.spyOn(isQuerySummaryCollapsed$, 'subscribe').mockImplementation((callback) => {
+      callback(isQuerySummaryCollapsed$.value);
+      return { unsubscribe: jest.fn() };
+    });
+    jest.spyOn(isSummaryAgentAvailable$, 'subscribe').mockImplementation((callback) => {
+      callback(isSummaryAgentAvailable$.value);
+      return { unsubscribe: jest.fn() };
+    });
+  });
+
+  afterEach(() => {
+    mockProps.resultSummaryEnabled$.complete();
+    mockProps.isQuerySummaryCollapsed$.complete();
+    mockProps.isSummaryAgentAvailable$.complete();
+  });
+
+  it('renders the context menu button with the correct label', () => {
+    const { getByText } = render(<ActionContextMenu {...mockProps} />);
+    expect(getByText('OpenSearch Assistant')).toBeInTheDocument();
+  });
+
+  it('should show the summarization action when there is Summary Agent', async () => {
+    render(<ActionContextMenu {...mockProps} />);
+    await act(async () => {
+      await mockAgentConfigExists.mock.results[0].value;
+    });
+    fireEvent.click(screen.getByLabelText('OpenSearch assistant trigger button'));
+    expect(screen.getByText('Show result summarization')).toBeInTheDocument();
+  });
+
+  it('should not show the summarization action when there is not Summary Agent', async () => {
+    const newIsSummaryAgentAvailable$ = new BehaviorSubject<boolean>(false);
+    const newMockAgentConfigExists = jest.fn().mockResolvedValue({ exists: false });
+    const newMockProps = {
+      ...mockProps,
+      isSummaryAgentAvailable$: newIsSummaryAgentAvailable$,
+      assistantServiceStart: {
+        ...mockProps.assistantServiceStart,
+        client: ({
+          ...mockProps.assistantServiceStart.client,
+          agentConfigExists: newMockAgentConfigExists,
+        } as unknown) as AssistantClient,
+      },
+    };
+    render(<ActionContextMenu {...newMockProps} />);
+
+    await act(async () => {
+      const resolvedValue = await newMockAgentConfigExists.mock.results[0].value;
+      newMockProps.isSummaryAgentAvailable$.next(resolvedValue.exists);
+    });
+
+    fireEvent.click(screen.getByLabelText('OpenSearch assistant trigger button'));
+
+    await waitFor(() => {
+      expect(screen.queryByTestId('queryAssist_summary_switch')).not.toBeInTheDocument();
+    });
+  });
+});

public/components/ui_action_context_menu.tsx

@@ -22,14 +22,15 @@ import { AI_ASSISTANT_QUERY_EDITOR_TRIGGER } from '../ui_triggers';
 import { getUiActions } from '../services';
 import { DataPublicPluginSetup } from '../../../../src/plugins/data/public';
 import { DATA2SUMMARY_AGENT_CONFIG_ID } from '../../common/constants/llm';
-import { checkAgentsExist } from '../../common/utils/llm_chat/get_is_data2summary_available';
-interface Props {
+import { AssistantServiceStart } from '../services/assistant_service';
+export interface Props {
   data: DataPublicPluginSetup;
   isQuerySummaryCollapsed$: BehaviorSubject<boolean>;
   resultSummaryEnabled$: BehaviorSubject<boolean>;
   isSummaryAgentAvailable$: BehaviorSubject<boolean>;
   httpSetup: HttpSetup;
   label?: string;
+  assistantServiceStart: AssistantServiceStart;
 }
 
 export const ActionContextMenu = (props: Props) => {
@@ -47,14 +48,13 @@ export const ActionContextMenu = (props: Props) => {
   const shouldShowSummarizationAction = resultSummaryEnabled && isSummaryAgentAvailable;
 
   useEffect(() => {
-    props.isSummaryAgentAvailable$.next(false);
     if (!resultSummaryEnabled) return;
+    props.isSummaryAgentAvailable$.next(false);
     const fetchSummaryAgent = async () => {
       try {
-        const summaryAgentStatus = await checkAgentsExist(
-          props.httpSetup,
+        const summaryAgentStatus = await props.assistantServiceStart.client.agentConfigExists(
           DATA2SUMMARY_AGENT_CONFIG_ID,
-          props.data.query.queryString.getQuery().dataset?.dataSource?.id
+          { dataSourceId: props.data.query.queryString.getQuery().dataset?.dataSource?.id }
         );
         props.isSummaryAgentAvailable$.next(!!summaryAgentStatus.exists);
       } catch (error) {

public/plugin.tsx

@@ -311,6 +311,7 @@ export class AssistantPlugin
                   resultSummaryEnabled$={resultSummaryEnabled$}
                   isSummaryAgentAvailable$={isSummaryAgentAvailable$}
                   data={setupDeps.data}
+                  assistantServiceStart={this.assistantService.start(core.http)}
                 />
               );
             },