Fix red underlined stuff in the IDE

Author: reakaleek

src/Elastic.Documentation.Site/Assets/web-components/SearchOrAskAi/AskAiAnswer.test.tsx

@@ -1,11 +1,15 @@
 import { AskAiAnswer } from './AskAiAnswer'
-import { render, screen, act } from '@testing-library/react'
+import { LlmGatewayMessage, useLlmGateway } from './useLlmGateway'
+import { render, screen } from '@testing-library/react'
 import userEvent from '@testing-library/user-event'
 import * as React from 'react'
+import { act } from 'react'
+
+const mockUseLlmGateway = jest.mocked(useLlmGateway)
 
-// Mock the dependencies
 const mockSendQuestion = jest.fn(() => Promise.resolve())
 const mockRetry = jest.fn()
+const mockAbort = jest.fn()
 
 jest.mock('./search.store', () => ({
     useAskAiTerm: jest.fn(() => 'What is Elasticsearch?'),
@@ -15,6 +19,7 @@ jest.mock('./useLlmGateway', () => ({
     useLlmGateway: jest.fn(() => ({
         messages: [],
         error: null,
+        abort: mockAbort,
         retry: mockRetry,
         sendQuestion: mockSendQuestion,
     })),
@@ -26,22 +31,19 @@ jest.mock('uuid', () => ({
 }))
 
 describe('AskAiAnswer Component', () => {
-    const {
-        useLlmGateway,
-    } = require('./useLlmGateway')
-
     beforeEach(() => {
         jest.clearAllMocks()
     })
 
     describe('Initial Loading State', () => {
         test('should show loading spinner when no messages are present', () => {
             // Arrange
-            useLlmGateway.mockReturnValue({
+            mockUseLlmGateway.mockReturnValue({
                 messages: [],
                 error: null,
                 retry: mockRetry,
                 sendQuestion: mockSendQuestion,
+                abort: mockAbort,
             })
 
             // Act
@@ -57,27 +59,32 @@ describe('AskAiAnswer Component', () => {
     describe('Message Display', () => {
         test('should display AI message content correctly', () => {
             // Arrange
-            const mockMessages = [
+            const mockMessages: LlmGatewayMessage[] = [
                 {
+                    id: 'some-id-1',
+                    timestamp: 0,
                     type: 'ai_message',
                     data: {
                         content:
                             'Elasticsearch is a distributed search engine...',
                     },
                 },
                 {
+                    id: 'some-id-2',
+                    timestamp: 0,
                     type: 'ai_message_chunk',
                     data: {
                         content: ' It provides real-time search capabilities.',
                     },
                 },
             ]
 
-            useLlmGateway.mockReturnValue({
+            mockUseLlmGateway.mockReturnValue({
                 messages: mockMessages,
                 error: null,
                 retry: mockRetry,
                 sendQuestion: mockSendQuestion,
+                abort: mockAbort,
             })
 
             // Act
@@ -93,11 +100,12 @@ describe('AskAiAnswer Component', () => {
     describe('Error State', () => {
         test('should display error message when there is an error', () => {
             // Arrange
-            useLlmGateway.mockReturnValue({
+            mockUseLlmGateway.mockReturnValue({
                 messages: [],
                 error: new Error('Network error'),
                 retry: mockRetry,
                 sendQuestion: mockSendQuestion,
+                abort: mockAbort,
             })
 
             // Act
@@ -118,24 +126,29 @@ describe('AskAiAnswer Component', () => {
     describe('Finished State with Feedback Buttons', () => {
         test('should show feedback buttons when answer is finished', () => {
             // Arrange
-            let onMessageCallback: (message: any) => void = () => {}
+            let onMessageCallback: (
+                message: LlmGatewayMessage
+            ) => void = () => {}
 
-            const mockMessages = [
+            const mockMessages: LlmGatewayMessage[] = [
                 {
+                    id: 'some-id-1',
+                    timestamp: 1,
                     type: 'ai_message',
                     data: {
                         content: 'Here is your answer about Elasticsearch.',
                     },
                 },
             ]
 
-            useLlmGateway.mockImplementation(({ onMessage }) => {
-                onMessageCallback = onMessage
+            mockUseLlmGateway.mockImplementation(({ onMessage }) => {
+                onMessageCallback = onMessage!
                 return {
                     messages: mockMessages,
                     error: null,
                     retry: mockRetry,
                     sendQuestion: mockSendQuestion,
+                    abort: mockAbort,
                 }
             })
 
@@ -144,7 +157,12 @@ describe('AskAiAnswer Component', () => {
 
             // Simulate the component receiving an 'agent_end' message to finish loading
             act(() => {
-                onMessageCallback({ type: 'agent_end' })
+                onMessageCallback({
+                    type: 'agent_end',
+                    id: 'some-id',
+                    timestamp: 12345,
+                    data: {},
+                })
             })
 
             // Assert
@@ -162,30 +180,46 @@ describe('AskAiAnswer Component', () => {
         test('should call retry function when refresh button is clicked', async () => {
             // Arrange
             const user = userEvent.setup()
-            let onMessageCallback: (message: any) => void = () => {}
+            let onMessageCallback: (
+                message: LlmGatewayMessage
+            ) => void = () => {}
 
-            const mockMessages = [
+            const mockMessages: LlmGatewayMessage[] = [
                 {
+                    id: 'some-id-1',
+                    timestamp: 12345,
                     type: 'ai_message',
                     data: { content: 'Here is your answer.' },
                 },
             ]
 
-            useLlmGateway.mockImplementation(({ onMessage }) => {
-                onMessageCallback = onMessage
+            mockUseLlmGateway.mockImplementation(({ onMessage }) => {
+                onMessageCallback = onMessage!
                 return {
                     messages: mockMessages,
                     error: null,
                     retry: mockRetry,
                     sendQuestion: mockSendQuestion,
+                    abort: mockAbort,
                 }
             })
 
             render(<AskAiAnswer />)
 
             // Simulate finished state
             act(() => {
-                onMessageCallback({ type: 'agent_end' })
+                onMessageCallback({
+                    type: 'agent_start',
+                    id: 'some-id',
+                    timestamp: 12345,
+                    data: { input: {}, thread: {} },
+                })
+                onMessageCallback({
+                    type: 'agent_end',
+                    id: 'some-id',
+                    timestamp: 12346,
+                    data: {},
+                })
             })
 
             // Act
@@ -203,11 +237,12 @@ describe('AskAiAnswer Component', () => {
     describe('Question Sending', () => {
         test('should send question on component mount', () => {
             // Arrange
-            useLlmGateway.mockReturnValue({
+            mockUseLlmGateway.mockReturnValue({
                 messages: [],
                 error: null,
                 retry: mockRetry,
                 sendQuestion: mockSendQuestion,
+                abort: mockAbort,
             })
 
             // Act

src/Elastic.Documentation.Site/tsconfig.json

@@ -14,8 +14,9 @@
     "isolatedModules": true,
     "noEmit": true,
     "jsx": "react-jsx",
-    "jsxImportSource": "@emotion/react"
+    "jsxImportSource": "@emotion/react",
+    "types": ["jest", "@testing-library/jest-dom"]
   },
-  "include": ["Assets/**/*"],
+  "include": ["Assets/**/*", "setupTests.ts"],
   "exclude": ["node_modules", "_static"]
 }