chore: fix test, combine preference to hook

Author: gagik

packages/compass-assistant/src/compass-assistant-provider.tsx

@@ -7,6 +7,7 @@ import { atlasServiceLocator } from '@mongodb-js/atlas-service/provider';
 import { DocsProviderTransport } from './docs-provider-transport';
 import { useDrawerActions } from '@mongodb-js/compass-components';
 import { buildExplainPlanPrompt } from './prompts';
+import { usePreference } from 'compass-preferences-model/provider';
 
 export const ASSISTANT_DRAWER_ID = 'compass-assistant-drawer';
 
@@ -31,8 +32,15 @@ export const AssistantActionsContext =
     interpretExplainPlan: () => {},
   });
 
-export function useAssistantActions(): AssistantActionsContextType {
-  return useContext(AssistantActionsContext);
+export function useAssistantActions(): AssistantActionsContextType & {
+  isAssistantEnabled: boolean;
+} {
+  const isAssistantEnabled = usePreference('enableAIAssistant');
+
+  return {
+    ...useContext(AssistantActionsContext),
+    isAssistantEnabled,
+  };
 }
 
 export const AssistantProvider: React.FunctionComponent<

packages/compass-explain-plan/src/components/explain-plan-modal.spec.tsx

@@ -21,7 +21,6 @@ function render(
     {
       dataService: {},
       localAppRegistry: {},
-      preferences: { enableAIAssistant: preferences?.enableAIAssistant },
     } as any,
     { on() {}, cleanup() {} } as any
   );
@@ -35,7 +34,8 @@ function render(
         onModalClose={() => {}}
         {...props}
       ></ExplainPlanModal>
-    </Provider>
+    </Provider>,
+    { preferences: { enableAIAssistant: preferences?.enableAIAssistant } }
   );
 }
 
@@ -70,4 +70,12 @@ describe('ExplainPlanModal', function () {
     render({ status: 'ready' }, { preferences: { enableAIAssistant: false } });
     expect(screen.queryByTestId('interpret-for-me-button')).to.not.exist;
   });
+
+  it('should disable the "Interpret for me" button when the status is not ready', function () {
+    render({ status: 'loading' }, { preferences: { enableAIAssistant: true } });
+    expect(screen.getByTestId('interpret-for-me-button')).to.have.attr(
+      'aria-disabled',
+      'true'
+    );
+  });
 });

packages/compass-explain-plan/src/components/explain-plan-modal.tsx

@@ -18,7 +18,6 @@ import { closeExplainPlanModal } from '../stores/explain-plan-modal-store';
 import { ExplainPlanView } from './explain-plan-view';
 import type { CollectionTabPluginMetadata } from '@mongodb-js/compass-collection';
 import { useAssistantActions } from '@mongodb-js/compass-assistant';
-import { usePreference } from 'compass-preferences-model/provider';
 
 export type ExplainPlanModalProps = Partial<
   Pick<
@@ -105,8 +104,7 @@ export const ExplainPlanModal: React.FunctionComponent<
   error,
   onModalClose,
 }) => {
-  const isAiAssistantEnabled = usePreference('enableAIAssistant');
-  const { interpretExplainPlan } = useAssistantActions();
+  const { interpretExplainPlan, isAssistantEnabled } = useAssistantActions();
 
   return (
     <Modal
@@ -134,7 +132,7 @@ export const ExplainPlanModal: React.FunctionComponent<
             }
           />
         </div>
-        {isAiAssistantEnabled && (
+        {isAssistantEnabled && (
           <div className={headerButtonSectionStyles}>
             <Button
               size="small"
@@ -148,6 +146,7 @@ export const ExplainPlanModal: React.FunctionComponent<
                 onModalClose();
                 interpretExplainPlan(JSON.stringify(explainPlan));
               }}
+              disabled={status !== 'ready'}
             >
               Interpret for me
             </Button>