From 203893f2503e6c50437da4c6476c3cf0e003a07b Mon Sep 17 00:00:00 2001 From: Nate Kelley Date: Fri, 9 May 2025 16:38:27 -0600 Subject: [PATCH] Unit tests and fix for replace message --- .../context/Chats/NewChatProvider.test.tsx | 316 ++++++++++++++++++ web/src/context/Chats/NewChatProvider.tsx | 6 +- 2 files changed, 318 insertions(+), 4 deletions(-) create mode 100644 web/src/context/Chats/NewChatProvider.test.tsx diff --git a/web/src/context/Chats/NewChatProvider.test.tsx b/web/src/context/Chats/NewChatProvider.test.tsx new file mode 100644 index 000000000..7ef802143 --- /dev/null +++ b/web/src/context/Chats/NewChatProvider.test.tsx @@ -0,0 +1,316 @@ +import { renderHook, act } from '@testing-library/react'; +import { useBusterNewChat } from './NewChatProvider'; +import { useBusterWebSocket } from '@/context/BusterWebSocket'; +import { useChatStreamMessage } from './useChatStreamMessage'; +import { useGetChatMemoized, useGetChatMessageMemoized } from '@/api/buster_rest/chats'; +import { useChatUpdate } from './useChatUpdate'; +import { create } from 'mutative'; +import { ShareAssetType } from '@/api/asset_interfaces'; + +// Mock dependencies +jest.mock('@/hooks', () => ({ + useMemoizedFn: (fn: any) => fn +})); + +jest.mock('@/context/BusterWebSocket'); +jest.mock('./useChatStreamMessage'); +jest.mock('@/api/buster_rest/chats'); +jest.mock('./useChatUpdate'); +jest.mock('mutative'); + +const mockUseBusterWebSocket = useBusterWebSocket as jest.Mock; +const mockUseChatStreamMessage = useChatStreamMessage as jest.Mock; +const mockUseGetChatMemoized = useGetChatMemoized as jest.Mock; +const mockUseGetChatMessageMemoized = useGetChatMessageMemoized as jest.Mock; +const mockUseChatUpdate = useChatUpdate as jest.Mock; +const mockCreate = create as jest.Mock; + +describe('useBusterNewChat', () => { + let mockBusterSocket: { + emitAndOnce: jest.Mock; + once: jest.Mock; + emit: jest.Mock; + }; + let mockInitializeNewChatCallback: jest.Mock; + let mockCompleteChatCallback: jest.Mock; + let mockStopChatCallback: jest.Mock; + let mockGetChatMemoized: jest.Mock; + let mockGetChatMessageMemoized: jest.Mock; + let mockOnUpdateChat: jest.Mock; + let mockOnUpdateChatMessage: jest.Mock; + + beforeEach(() => { + mockBusterSocket = { + emitAndOnce: jest.fn().mockResolvedValue({}), + once: jest.fn(), + emit: jest.fn() + }; + mockUseBusterWebSocket.mockReturnValue(mockBusterSocket); + + mockInitializeNewChatCallback = jest.fn(); + mockCompleteChatCallback = jest.fn(); + mockStopChatCallback = jest.fn(); + mockUseChatStreamMessage.mockReturnValue({ + initializeNewChatCallback: mockInitializeNewChatCallback, + completeChatCallback: mockCompleteChatCallback, + stopChatCallback: mockStopChatCallback + }); + + mockGetChatMemoized = jest.fn(); + mockUseGetChatMemoized.mockReturnValue(mockGetChatMemoized); + + mockGetChatMessageMemoized = jest.fn(); + mockUseGetChatMessageMemoized.mockReturnValue(mockGetChatMessageMemoized); + + mockOnUpdateChat = jest.fn(); + mockOnUpdateChatMessage = jest.fn(); + mockUseChatUpdate.mockReturnValue({ + onUpdateChat: mockOnUpdateChat, + onUpdateChatMessage: mockOnUpdateChatMessage + }); + + mockCreate.mockImplementation((base, updater) => { + const draft = JSON.parse(JSON.stringify(base)); + updater(draft); + return draft; + }); + }); + + afterEach(() => { + jest.clearAllMocks(); + }); + + // Test 1: onSelectSearchAsset should resolve after a delay (mocked) + test('onSelectSearchAsset should complete', async () => { + jest.useFakeTimers(); + const { result } = renderHook(() => useBusterNewChat()); + const promise = result.current.onSelectSearchAsset({ + id: 'asset1', + name: 'Asset 1', + type: ShareAssetType.METRIC, + highlights: [], + updated_at: new Date().toISOString(), + score: 0 + }); + + act(() => { + jest.runAllTimers(); + }); + + await expect(promise).resolves.toBeUndefined(); + jest.useRealTimers(); + }); + + // Test 2: onStartNewChat should call busterSocket.emitAndOnce and busterSocket.once + test('onStartNewChat should call socket methods with correct parameters', async () => { + const { result } = renderHook(() => useBusterNewChat()); + const chatPayload = { prompt: 'Hello' }; + + await result.current.onStartNewChat(chatPayload); + + expect(mockBusterSocket.emitAndOnce).toHaveBeenCalledWith({ + emitEvent: { + route: '/chats/post', + payload: { + dataset_id: undefined, + prompt: 'Hello', + metric_id: undefined, + dashboard_id: undefined + } + }, + responseEvent: { + route: '/chats/post:initializeChat', + callback: mockInitializeNewChatCallback + } + }); + expect(mockBusterSocket.once).toHaveBeenCalledWith({ + route: '/chats/post:complete', + callback: mockCompleteChatCallback + }); + }); + + // Test 3: onStartNewChat should include datasetId if provided + test('onStartNewChat should include datasetId when provided', async () => { + const { result } = renderHook(() => useBusterNewChat()); + const chatPayload = { prompt: 'Test with dataset', datasetId: 'ds1' }; + + await result.current.onStartNewChat(chatPayload); + + expect(mockBusterSocket.emitAndOnce).toHaveBeenCalledWith( + expect.objectContaining({ + emitEvent: expect.objectContaining({ + payload: expect.objectContaining({ dataset_id: 'ds1' }) + }) + }) + ); + }); + + // Test 4: onStartChatFromFile should call onStartNewChat with metricId for metric fileType + test('onStartChatFromFile should call onStartNewChat with metricId for metric type', async () => { + const { result } = renderHook(() => useBusterNewChat()); + + const filePayload = { + prompt: 'Chat from metric', + fileId: 'metric123', + fileType: 'metric' as const + }; + await result.current.onStartChatFromFile(filePayload); + + expect(mockBusterSocket.emitAndOnce).toHaveBeenCalledWith( + expect.objectContaining({ + emitEvent: expect.objectContaining({ + payload: expect.objectContaining({ + prompt: 'Chat from metric', + metric_id: 'metric123', + dashboard_id: undefined + }) + }) + }) + ); + }); + + // Test 5: onStartChatFromFile should call onStartNewChat with dashboardId for dashboard fileType + test('onStartChatFromFile should call onStartNewChat with dashboardId for dashboard type', async () => { + const { result } = renderHook(() => useBusterNewChat()); + + const filePayload = { + prompt: 'Chat from dashboard', + fileId: 'dash123', + fileType: 'dashboard' as const + }; + await result.current.onStartChatFromFile(filePayload); + + expect(mockBusterSocket.emitAndOnce).toHaveBeenCalledWith( + expect.objectContaining({ + emitEvent: expect.objectContaining({ + payload: expect.objectContaining({ + prompt: 'Chat from dashboard', + metric_id: undefined, + dashboard_id: 'dash123' + }) + }) + }) + ); + }); + + // Test 6: onFollowUpChat should call socket methods with correct parameters + test('onFollowUpChat should call socket methods with chatId', async () => { + const { result } = renderHook(() => useBusterNewChat()); + const followUpPayload = { prompt: 'Follow up question', chatId: 'chat1' }; + + await result.current.onFollowUpChat(followUpPayload); + + expect(mockBusterSocket.once).toHaveBeenCalledWith({ + route: '/chats/post:initializeChat', + callback: mockInitializeNewChatCallback + }); + expect(mockBusterSocket.emitAndOnce).toHaveBeenCalledWith({ + emitEvent: { + route: '/chats/post', + payload: { + prompt: 'Follow up question', + chat_id: 'chat1' + } + }, + responseEvent: { + route: '/chats/post:complete', + callback: mockCompleteChatCallback + } + }); + }); + + // Test 7: onStopChat should call busterSocket.emit and stopChatCallback + test('onStopChat should call emit and stopChatCallback', () => { + const { result } = renderHook(() => useBusterNewChat()); + const stopPayload = { chatId: 'chat1', messageId: 'msg1' }; + + result.current.onStopChat(stopPayload); + + expect(mockBusterSocket.emit).toHaveBeenCalledWith({ + route: '/chats/stop', + payload: { + id: 'chat1', + message_id: 'msg1' + } + }); + expect(mockStopChatCallback).toHaveBeenCalledWith('chat1'); + }); + + // Test 8: onReplaceMessageInChat updates message and chat correctly and calls socket + test('onReplaceMessageInChat should update message, chat, and call socket', async () => { + const mockCurrentChat = { id: 'chat1', message_ids: ['msg0', 'msg1', 'msg2'] }; + const mockCurrentMessage = { id: 'msg1', request_message: { request: 'Old prompt' } }; + + mockGetChatMemoized.mockReturnValue(mockCurrentChat); + mockGetChatMessageMemoized.mockReturnValue(mockCurrentMessage); + + const { result } = renderHook(() => useBusterNewChat()); + const replacePayload = { prompt: 'New prompt', messageId: 'msg1', chatId: 'chat1' }; + + await result.current.onReplaceMessageInChat(replacePayload); + + expect(mockGetChatMemoized).toHaveBeenCalledWith('chat1'); + expect(mockGetChatMessageMemoized).toHaveBeenCalledWith('msg1'); + + expect(mockOnUpdateChatMessage).toHaveBeenCalledWith({ + id: 'msg1', + request_message: { request: 'New prompt' }, + reasoning_message_ids: [], + response_message_ids: [], + reasoning_messages: {}, + final_reasoning_message: null, + isCompletedStream: false + }); + + expect(mockOnUpdateChat).toHaveBeenCalledWith({ + id: 'chat1', + message_ids: ['msg0', 'msg1'] + }); + + expect(mockBusterSocket.emitAndOnce).toHaveBeenCalledWith({ + emitEvent: { + route: '/chats/post', + payload: { + prompt: 'New prompt', + message_id: 'msg1', + chat_id: 'chat1' + } + }, + responseEvent: { + route: '/chats/post:complete', + callback: mockCompleteChatCallback + } + }); + }); + + // Test 9: onReplaceMessageInChat handles message not found in chat but still calls socket + test('onReplaceMessageInChat should not update chat if message not found, but still calls socket', async () => { + const mockCurrentChat = { id: 'chat1', message_ids: ['msg0', 'msg2'] }; // msg1 is missing + const mockCurrentMessage = { id: 'msg1', request_message: { request: 'Old prompt' } }; + + mockGetChatMemoized.mockReturnValue(mockCurrentChat); + mockGetChatMessageMemoized.mockReturnValue(mockCurrentMessage); + + const { result } = renderHook(() => useBusterNewChat()); + const replacePayload = { prompt: 'New prompt', messageId: 'msg1', chatId: 'chat1' }; + + await result.current.onReplaceMessageInChat(replacePayload); + + expect(mockOnUpdateChatMessage).toHaveBeenCalled(); + expect(mockOnUpdateChat).not.toHaveBeenCalled(); + expect(mockBusterSocket.emitAndOnce).toHaveBeenCalledWith({ + emitEvent: { + route: '/chats/post', + payload: { + prompt: 'New prompt', + message_id: 'msg1', + chat_id: 'chat1' + } + }, + responseEvent: { + route: '/chats/post:complete', + callback: mockCompleteChatCallback + } + }); + }); +}); diff --git a/web/src/context/Chats/NewChatProvider.tsx b/web/src/context/Chats/NewChatProvider.tsx index df186b971..825996b13 100644 --- a/web/src/context/Chats/NewChatProvider.tsx +++ b/web/src/context/Chats/NewChatProvider.tsx @@ -89,6 +89,8 @@ export const useBusterNewChat = () => { const currentChat = getChatMemoized(chatId); const currentMessage = getChatMessageMemoized(messageId); const currentRequestMessage = currentMessage?.request_message!; + const messageIndex = currentChat?.message_ids.findIndex((mId) => mId === messageId); + onUpdateChatMessage({ id: messageId, request_message: create(currentRequestMessage, (draft) => { @@ -101,10 +103,6 @@ export const useBusterNewChat = () => { isCompletedStream: false }); - const messageIndex = currentChat?.message_ids.findIndex( - (messageId) => messageId === messageId - ); - if (messageIndex !== -1 && typeof messageIndex === 'number') { const updatedMessageIds = currentChat?.message_ids.slice(0, messageIndex + 1); onUpdateChat({