[9.1] add loading state install/uninstall product doc (#229579)

x-pack/platform/plugins/private/observability_ai_assistant_management/public/hooks/use_product_doc.ts

@@ -0,0 +1,95 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { useEffect, useMemo } from 'react';
+import { useQuery } from '@tanstack/react-query';
+import { InstallationStatus } from '@kbn/product-doc-base-plugin/common/install_status';
+import {
+  PerformInstallResponse,
+  UninstallResponse,
+} from '@kbn/product-doc-base-plugin/common/http_api/installation';
+import { REACT_QUERY_KEYS } from '../constants';
+import { useKibana } from './use_kibana';
+import { useUninstallProductDoc } from './use_uninstall_product_doc';
+import { useInstallProductDoc } from './use_install_product_doc';
+
+export interface UseProductDoc {
+  status: InstallationStatus | undefined;
+  isLoading: boolean;
+  installProductDoc: (inferenceId: string) => Promise<PerformInstallResponse>;
+  uninstallProductDoc: (inferenceId: string) => Promise<UninstallResponse>;
+}
+
+/**
+ * Custom hook to get the status of the product documentation installation.
+ * It also provides methods to install and uninstall the product documentation.
+ *
+ * @param inferenceId - The ID of the inference for which to get the product documentation status.
+ * @returns An object containing the status of the product documentation, loading state, and methods to install and uninstall the product documentation.
+ */
+export function useProductDoc(inferenceId: string | undefined): UseProductDoc {
+  const { productDocBase } = useKibana().services;
+
+  const { mutateAsync: installProductDoc, isLoading: isInstalling } = useInstallProductDoc();
+
+  const { mutateAsync: uninstallProductDoc, isLoading: isUninstalling } = useUninstallProductDoc();
+
+  const { isLoading, data, refetch, isRefetching } = useQuery({
+    networkMode: 'always',
+    queryKey: [REACT_QUERY_KEYS.GET_PRODUCT_DOC_STATUS, inferenceId],
+    queryFn: async () => {
+      return productDocBase!.installation.getStatus({ inferenceId });
+    },
+    keepPreviousData: false,
+    refetchOnWindowFocus: false,
+  });
+
+  useEffect(() => {
+    refetch();
+  }, [inferenceId, refetch]);
+
+  // poll the status if when is installing or uninstalling
+  useEffect(() => {
+    if (
+      !(
+        data?.overall === 'installing' ||
+        data?.overall === 'uninstalling' ||
+        isInstalling ||
+        isUninstalling
+      )
+    ) {
+      return;
+    }
+
+    const interval = setInterval(refetch, 5000);
+
+    // cleanup the interval if unmount
+    return () => {
+      clearInterval(interval);
+    };
+  }, [refetch, data?.overall, isInstalling, isUninstalling]);
+
+  const status: InstallationStatus | undefined = useMemo(() => {
+    if (!inferenceId || data?.inferenceId !== inferenceId) {
+      return undefined;
+    }
+    if (isInstalling) {
+      return 'installing';
+    }
+    if (isUninstalling) {
+      return 'uninstalling';
+    }
+    return data?.overall;
+  }, [inferenceId, isInstalling, isUninstalling, data]);
+
+  return {
+    status,
+    isLoading: isLoading || isRefetching || isInstalling || isUninstalling,
+    installProductDoc,
+    uninstallProductDoc,
+  };
+}

x-pack/platform/plugins/private/observability_ai_assistant_management/public/routes/components/settings_tab/change_kb_model.test.tsx

@@ -24,6 +24,8 @@ import {
   elserTitle,
 } from '@kbn/ai-assistant/src/utils/get_model_options_for_inference_endpoints';
 import { QueryClient, QueryClientProvider } from '@tanstack/react-query';
+import { InstallationStatus } from '@kbn/product-doc-base-plugin/common/install_status';
+import { UseProductDoc } from '../../../hooks/use_product_doc';
 
 jest.mock('../../../hooks/use_install_product_doc', () => ({
   useInstallProductDoc: () => ({
@@ -94,6 +96,14 @@ const createMockKnowledgeBase = (
   ...overrides,
 });
 
+const createProductDoc = (overrides: Partial<UseProductDoc> = {}) => ({
+  status: 'uninstalled' as InstallationStatus,
+  isLoading: false,
+  installProductDoc: jest.fn().mockResolvedValue({} as any),
+  uninstallProductDoc: jest.fn().mockResolvedValue({} as any),
+  ...overrides,
+});
+
 const modelOptions = [
   {
     key: ELSER_ON_ML_NODE_INFERENCE_ID,
@@ -112,12 +122,16 @@ const setupMockGetModelOptions = (options = modelOptions) => {
   mockGetModelOptions.mockReturnValue(options);
 };
 
-const renderComponent = (mockKb: UseKnowledgeBaseResult) => {
+const renderComponent = (mockKb: UseKnowledgeBaseResult, mockProductDoc: UseProductDoc) => {
   const queryClient = new QueryClient();
 
   render(
     <QueryClientProvider client={queryClient}>
-      <ChangeKbModel knowledgeBase={mockKb} />{' '}
+      <ChangeKbModel
+        knowledgeBase={mockKb}
+        productDoc={mockProductDoc}
+        currentlyDeployedInferenceId={ELSER_ON_ML_NODE_INFERENCE_ID}
+      />
     </QueryClientProvider>
   );
 };
@@ -133,15 +147,17 @@ describe('ChangeKbModel', () => {
 
   it('disables the `Update` button when selected model is the same as current and no redeployment needed', () => {
     const mockKb = createMockKnowledgeBase();
-    renderComponent(mockKb);
+    const mockProductDoc = createProductDoc();
+    renderComponent(mockKb, mockProductDoc);
 
     const button = screen.getByTestId('observabilityAiAssistantKnowledgeBaseUpdateModelButton');
     expect(button).toBeDisabled();
   });
 
   it('enables the `Update` button when a different model is selected', async () => {
     const mockKb = createMockKnowledgeBase();
-    renderComponent(mockKb);
+    const mockProductDoc = createProductDoc();
+    renderComponent(mockKb, mockProductDoc);
 
     const button = screen.getByTestId('observabilityAiAssistantKnowledgeBaseUpdateModelButton');
     expect(button).toBeDisabled();
@@ -159,7 +175,8 @@ describe('ChangeKbModel', () => {
 
   it('disables the `Update` button when knowledge base is installing', () => {
     const mockKb = createMockKnowledgeBase({ isInstalling: true });
-    renderComponent(mockKb);
+    const mockProductDoc = createProductDoc();
+    renderComponent(mockKb, mockProductDoc);
 
     const button = screen.getByTestId('observabilityAiAssistantKnowledgeBaseUpdateModelButton');
     expect(button).toBeDisabled();
@@ -180,7 +197,8 @@ describe('ChangeKbModel', () => {
           },
         }),
       });
-      renderComponent(mockKb);
+      const mockProductDoc = createProductDoc();
+      renderComponent(mockKb, mockProductDoc);
 
       const dropdown = screen.getByTestId('observabilityAiAssistantKnowledgeBaseModelDropdown');
       expect(dropdown).toHaveTextContent(elserTitle);
@@ -198,7 +216,8 @@ describe('ChangeKbModel', () => {
           },
         }),
       });
-      renderComponent(mockKb);
+      const mockProductDoc = createProductDoc();
+      renderComponent(mockKb, mockProductDoc);
 
       const dropdown = screen.getByTestId('observabilityAiAssistantKnowledgeBaseModelDropdown');
       dropdown.click();
@@ -226,7 +245,8 @@ describe('ChangeKbModel', () => {
           },
         }),
       });
-      renderComponent(mockKb);
+      const mockProductDoc = createProductDoc();
+      renderComponent(mockKb, mockProductDoc);
 
       const dropdown = screen.getByTestId('observabilityAiAssistantKnowledgeBaseModelDropdown');
       dropdown.click();

x-pack/platform/plugins/private/observability_ai_assistant_management/public/routes/components/settings_tab/change_kb_model.tsx

@@ -31,10 +31,17 @@ import {
   LEGACY_CUSTOM_INFERENCE_ID,
   useKibana,
 } from '@kbn/observability-ai-assistant-plugin/public';
-import { getMappedInferenceId } from '../../../helpers/inference_utils';
-import { useGetProductDoc } from '../../../hooks/use_get_product_doc';
+import { UseProductDoc } from '../../../hooks/use_product_doc';
 
-export function ChangeKbModel({ knowledgeBase }: { knowledgeBase: UseKnowledgeBaseResult }) {
+export function ChangeKbModel({
+  knowledgeBase,
+  productDoc,
+  currentlyDeployedInferenceId,
+}: {
+  knowledgeBase: UseKnowledgeBaseResult;
+  productDoc: UseProductDoc;
+  currentlyDeployedInferenceId: string | undefined;
+}) {
   const { overlays } = useKibana().services;
 
   const [hasLoadedCurrentModel, setHasLoadedCurrentModel] = useState(false);
@@ -46,12 +53,6 @@ export function ChangeKbModel({ knowledgeBase }: { knowledgeBase: UseKnowledgeBa
     endpoints: inferenceEndpoints,
   });
 
-  const currentlyDeployedInferenceId = getMappedInferenceId(
-    knowledgeBase.status.value?.currentInferenceId
-  );
-
-  const { installProductDoc } = useGetProductDoc(currentlyDeployedInferenceId);
-
   const [selectedInferenceId, setSelectedInferenceId] = useState(
     currentlyDeployedInferenceId || ''
   );
@@ -165,7 +166,7 @@ export function ChangeKbModel({ knowledgeBase }: { knowledgeBase: UseKnowledgeBa
             if (isConfirmed) {
               setIsUpdatingModel(true);
               knowledgeBase.install(selectedInferenceId);
-              installProductDoc(selectedInferenceId);
+              productDoc.installProductDoc(selectedInferenceId);
             }
           });
       }
@@ -177,7 +178,7 @@ export function ChangeKbModel({ knowledgeBase }: { knowledgeBase: UseKnowledgeBa
     isSelectedModelCurrentModel,
     overlays,
     confirmationMessages,
-    installProductDoc,
+    productDoc,
   ]);
 
   const superSelectOptions = modelOptions.map((option: ModelOptionsData) => ({

x-pack/platform/plugins/private/observability_ai_assistant_management/public/routes/components/settings_tab/product_doc_entry.test.tsx

@@ -15,19 +15,8 @@ import {
   LEGACY_CUSTOM_INFERENCE_ID,
 } from '@kbn/observability-ai-assistant-plugin/public';
 import { UseKnowledgeBaseResult } from '@kbn/ai-assistant';
-import { useGetProductDoc } from '../../../hooks/use_get_product_doc';
-
-jest.mock('../../../hooks/use_get_product_doc', () => ({
-  useGetProductDoc: jest.fn(),
-}));
-
-jest.mock('../../../hooks/use_install_product_doc', () => ({
-  useInstallProductDoc: () => ({ mutateAsync: jest.fn() }),
-}));
-
-jest.mock('../../../hooks/use_uninstall_product_doc', () => ({
-  useUninstallProductDoc: () => ({ mutateAsync: jest.fn() }),
-}));
+import { UseProductDoc } from '../../../hooks/use_product_doc';
+import { InstallationStatus } from '@kbn/product-doc-base-plugin/common/install_status';
 
 const createMockStatus = (
   overrides?: Partial<APIReturnType<'GET /internal/observability_ai_assistant/kb/status'>>
@@ -62,12 +51,16 @@ const createMockKnowledgeBase = (
   ...overrides,
 });
 
-describe('ProductDocEntry', () => {
-  it('calls useGetProductDocStatus with ELSER_ON_ML_NODE_INFERENCE_ID when inference ID is LEGACY_CUSTOM_INFERENCE_ID', async () => {
-    (useGetProductDoc as jest.Mock).mockReturnValue({
-      status: 'installed',
-    });
+const createProductDoc = (overrides: Partial<UseProductDoc> = {}) => ({
+  status: 'uninstalled' as InstallationStatus,
+  isLoading: false,
+  installProductDoc: jest.fn().mockResolvedValue({} as any),
+  uninstallProductDoc: jest.fn().mockResolvedValue({} as any),
+  ...overrides,
+});
 
+describe('ProductDocEntry', () => {
+  it('should render the installed state correctly', async () => {
     const mockKnowledgeBase = createMockKnowledgeBase({
       status: createMockStatus({
         currentInferenceId: LEGACY_CUSTOM_INFERENCE_ID,
@@ -80,28 +73,37 @@ describe('ProductDocEntry', () => {
       }),
     });
 
-    render(<ProductDocEntry knowledgeBase={mockKnowledgeBase} />);
+    const productDoc = createProductDoc({
+      status: 'installed',
+    });
+
+    render(
+      <ProductDocEntry
+        knowledgeBase={mockKnowledgeBase}
+        productDoc={productDoc}
+        currentlyDeployedInferenceId={undefined}
+      />
+    );
 
     await waitFor(() => {
       expect(screen.getByText('Installed')).toBeInTheDocument();
     });
-
-    expect(useGetProductDoc).toHaveBeenCalledWith(ELSER_ON_ML_NODE_INFERENCE_ID);
   });
 
-  it('calls useGetProductDocStatus with the current inference ID when inference ID is not LEGACY_CUSTOM_INFERENCE_ID"', async () => {
+  it('should render the uninstalled state correctly', async () => {
     const mockKnowledgeBase = createMockKnowledgeBase();
+    const productDoc = createProductDoc();
 
-    (useGetProductDoc as jest.Mock).mockReturnValue({
-      status: 'installed',
-    });
-
-    render(<ProductDocEntry knowledgeBase={mockKnowledgeBase} />);
+    render(
+      <ProductDocEntry
+        knowledgeBase={mockKnowledgeBase}
+        productDoc={productDoc}
+        currentlyDeployedInferenceId={undefined}
+      />
+    );
 
     await waitFor(() => {
-      expect(screen.getByText('Installed')).toBeInTheDocument();
+      expect(screen.getByText('Install')).toBeInTheDocument();
     });
-
-    expect(useGetProductDoc).toHaveBeenCalledWith(ELSER_ON_ML_NODE_INFERENCE_ID);
   });
 });