[Security solution] Assistant race condition bug fixing (#187186)

This commit is contained in:
Steph Milovic 2024-07-02 13:42:22 -06:00 committed by GitHub
parent 2aa94a27f0
commit ad4fe84078
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 39 additions and 28 deletions

View file

@ -125,6 +125,7 @@ export const AssistantHeaderFlyout: React.FC<Props> = ({
`,
onClick: showDestroyModal,
icon: 'refresh',
'data-test-subj': 'clear-chat',
},
],
},
@ -243,6 +244,7 @@ export const AssistantHeaderFlyout: React.FC<Props> = ({
aria-label="test"
iconType="boxesVertical"
onClick={onButtonClick}
data-test-subj="chat-context-menu"
/>
}
isOpen={isPopoverOpen}
@ -266,6 +268,7 @@ export const AssistantHeaderFlyout: React.FC<Props> = ({
confirmButtonText={i18n.RESET_BUTTON_TEXT}
buttonColor="danger"
defaultFocusedButton="confirm"
data-test-subj="reset-conversation-modal"
>
<p>{i18n.CLEAR_CHAT_CONFIRMATION}</p>
</EuiConfirmModal>

View file

@ -35,7 +35,7 @@ export interface UseChatSendProps {
export interface UseChatSend {
abortStream: () => void;
handleOnChatCleared: () => void;
handleOnChatCleared: () => Promise<void>;
handlePromptChange: (prompt: string) => void;
handleSendMessage: (promptText: string) => void;
handleRegenerateResponse: () => void;

View file

@ -7,7 +7,7 @@
import React from 'react';
import { act, fireEvent, render, screen, within } from '@testing-library/react';
import { act, fireEvent, render, screen, waitFor, within } from '@testing-library/react';
import { Assistant } from '.';
import type { IHttpFetchError } from '@kbn/core/public';
@ -63,15 +63,25 @@ const mockData = {
},
};
const mockDeleteConvo = jest.fn();
const clearConversation = jest.fn();
const mockUseConversation = {
clearConversation: clearConversation.mockResolvedValue(mockData.welcome_id),
getConversation: jest.fn(),
getDefaultConversation: jest.fn().mockReturnValue(mockData.welcome_id),
deleteConversation: mockDeleteConvo,
setApiConfig: jest.fn().mockResolvedValue({}),
};
const refetchResults = jest.fn();
describe('Assistant', () => {
beforeAll(() => {
let persistToLocalStorage: jest.Mock;
let persistToSessionStorage: jest.Mock;
beforeEach(() => {
jest.clearAllMocks();
persistToLocalStorage = jest.fn();
persistToSessionStorage = jest.fn();
(useConversation as jest.Mock).mockReturnValue(mockUseConversation);
jest.mocked(useConnectorSetup).mockReturnValue({
comments: [],
@ -89,13 +99,14 @@ describe('Assistant', () => {
];
jest.mocked(useLoadConnectors).mockReturnValue({
isFetched: true,
isFetchedAfterMount: true,
data: connectors,
} as unknown as UseQueryResult<AIConnector[], IHttpFetchError>);
jest.mocked(useFetchCurrentUserConversations).mockReturnValue({
data: mockData,
isLoading: false,
refetch: jest.fn().mockResolvedValue({
refetch: refetchResults.mockResolvedValue({
isLoading: false,
data: {
...mockData,
@ -107,16 +118,6 @@ describe('Assistant', () => {
}),
isFetched: true,
} as unknown as DefinedUseQueryResult<Record<string, Conversation>, unknown>);
});
let persistToLocalStorage: jest.Mock;
let persistToSessionStorage: jest.Mock;
beforeEach(() => {
jest.clearAllMocks();
persistToLocalStorage = jest.fn();
persistToSessionStorage = jest.fn();
jest
.mocked(useLocalStorage)
.mockReturnValue([undefined, persistToLocalStorage] as unknown as ReturnType<
@ -234,6 +235,16 @@ describe('Assistant', () => {
});
expect(mockDeleteConvo).toHaveBeenCalledWith(mockData.welcome_id.id);
});
it('should refetchConversationsState after clear chat history button click', async () => {
renderAssistant({ isFlyoutMode: true });
fireEvent.click(screen.getByTestId('chat-context-menu'));
fireEvent.click(screen.getByTestId('clear-chat'));
fireEvent.click(screen.getByTestId('confirmModalConfirmButton'));
await waitFor(() => {
expect(clearConversation).toHaveBeenCalled();
expect(refetchResults).toHaveBeenCalled();
});
});
});
describe('when selected conversation changes and some connectors are loaded', () => {
it('should persist the conversation id to local storage', async () => {

View file

@ -220,7 +220,7 @@ const AssistantComponent: React.FC<Props> = ({
);
useEffect(() => {
if (conversationsLoaded && Object.keys(conversations).length > 0) {
if (areConnectorsFetched && conversationsLoaded && Object.keys(conversations).length > 0) {
setCurrentConversation((prev) => {
const nextConversation =
(currentConversationId && conversations[currentConversationId]) ||
@ -256,13 +256,13 @@ const AssistantComponent: React.FC<Props> = ({
});
}
}, [
areConnectorsFetched,
conversationTitle,
conversations,
conversationsLoaded,
currentConversationId,
getDefaultConversation,
getLastConversationId,
conversationsLoaded,
currentConversation?.id,
currentConversationId,
isAssistantEnabled,
isFlyoutMode,
]);
@ -549,7 +549,7 @@ const AssistantComponent: React.FC<Props> = ({
const {
abortStream,
handleOnChatCleared,
handleOnChatCleared: onChatCleared,
handlePromptChange,
handleSendMessage,
handleRegenerateResponse,
@ -567,6 +567,11 @@ const AssistantComponent: React.FC<Props> = ({
setCurrentConversation,
});
const handleOnChatCleared = useCallback(async () => {
await onChatCleared();
await refetchResults();
}, [onChatCleared, refetchResults]);
const handleChatSend = useCallback(
async (promptText: string) => {
await handleSendMessage(promptText);
@ -733,15 +738,7 @@ const AssistantComponent: React.FC<Props> = ({
}
}
})();
}, [
currentConversation,
defaultConnector,
refetchConversationsState,
setApiConfig,
showMissingConnectorCallout,
areConnectorsFetched,
mutateAsync,
]);
}, [areConnectorsFetched, currentConversation, mutateAsync]);
const handleCreateConversation = useCallback(async () => {
const newChatExists = find(conversations, ['title', NEW_CHAT]);