test fix

elastic · Oct 30, 2024 · c55e216 · c55e216
1 parent 1bd64de
commit c55e216
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 8 deletions.
diff --git a/x-pack/plugins/elastic_assistant/server/routes/anonymization_fields/find_route.test.ts b/x-pack/plugins/elastic_assistant/server/routes/anonymization_fields/find_route.test.ts
@@ -12,6 +12,7 @@ import { requestContextMock } from '../../__mocks__/request_context';
 import { getFindAnonymizationFieldsResultWithSingleHit } from '../../__mocks__/response';
 import { findAnonymizationFieldsRoute } from './find_route';
 import { loggingSystemMock } from '@kbn/core-logging-server-mocks';
+import type { AuthenticatedUser } from '@kbn/core-security-common';
 
 describe('Find user anonymization fields route', () => {
   let server: ReturnType<typeof serverMock.create>;
@@ -25,13 +26,13 @@ describe('Find user anonymization fields route', () => {
     clients.elasticAssistant.getAIAssistantAnonymizationFieldsDataClient.findDocuments.mockResolvedValue(
       Promise.resolve(getFindAnonymizationFieldsResultWithSingleHit())
     );
-    clients.elasticAssistant.getCurrentUser.mockResolvedValue({
+    context.elasticAssistant.getCurrentUser.mockReturnValue({
       username: 'my_username',
       authentication_realm: {
         type: 'my_realm_type',
         name: 'my_realm_name',
       },
-    });
+    } as AuthenticatedUser);
     logger = loggingSystemMock.createLogger();
 
     findAnonymizationFieldsRoute(server.router, logger);

diff --git a/x-pack/plugins/elastic_assistant/server/routes/chat/chat_complete_route.test.ts b/x-pack/plugins/elastic_assistant/server/routes/chat/chat_complete_route.test.ts
@@ -32,7 +32,17 @@ const actionsClient = actionsClientMock.create();
 jest.mock('../../lib/build_response', () => ({
   buildResponse: jest.fn().mockImplementation((x) => x),
 }));
-jest.mock('../helpers');
+
+jest.mock('../helpers', () => {
+  const original = jest.requireActual('../helpers');
+
+  return {
+    ...original,
+    appendAssistantMessageToConversation: jest.fn(),
+    createConversationWithUserInput: jest.fn(),
+    langChainExecute: jest.fn(),
+  };
+});
 const mockAppendAssistantMessageToConversation = appendAssistantMessageToConversation as jest.Mock;
 
 const mockLangChainExecute = langChainExecute as jest.Mock;

diff --git a/x-pack/plugins/elastic_assistant/server/routes/prompts/find_route.test.ts b/x-pack/plugins/elastic_assistant/server/routes/prompts/find_route.test.ts
@@ -12,6 +12,7 @@ import { requestContextMock } from '../../__mocks__/request_context';
 import { getFindPromptsResultWithSingleHit } from '../../__mocks__/response';
 import { findPromptsRoute } from './find_route';
 import { loggingSystemMock } from '@kbn/core-logging-server-mocks';
+import type { AuthenticatedUser } from '@kbn/core-security-common';
 
 describe('Find user prompts route', () => {
   let server: ReturnType<typeof serverMock.create>;
@@ -25,13 +26,13 @@ describe('Find user prompts route', () => {
     clients.elasticAssistant.getAIAssistantPromptsDataClient.findDocuments.mockResolvedValue(
       Promise.resolve(getFindPromptsResultWithSingleHit())
     );
-    clients.elasticAssistant.getCurrentUser.mockResolvedValue({
+    context.elasticAssistant.getCurrentUser.mockReturnValue({
       username: 'my_username',
       authentication_realm: {
         type: 'my_realm_type',
         name: 'my_realm_name',
       },
-    });
+    } as AuthenticatedUser);
     logger = loggingSystemMock.createLogger();
 
     findPromptsRoute(server.router, logger);

diff --git a/x-pack/plugins/elastic_assistant/server/routes/user_conversations/find_route.test.ts b/x-pack/plugins/elastic_assistant/server/routes/user_conversations/find_route.test.ts
@@ -4,7 +4,7 @@
  * 2.0; you may not use this file except in compliance with the Elastic License
  * 2.0.
  */
-
+import { type AuthenticatedUser } from '@kbn/core/server';
 import { getCurrentUserFindRequest, requestMock } from '../../__mocks__/request';
 import { ELASTIC_AI_ASSISTANT_CONVERSATIONS_URL_FIND } from '@kbn/elastic-assistant-common';
 import { serverMock } from '../../__mocks__/server';
@@ -23,13 +23,13 @@ describe('Find user conversations route', () => {
     clients.elasticAssistant.getAIAssistantConversationsDataClient.findDocuments.mockResolvedValue(
       Promise.resolve(getFindConversationsResultWithSingleHit())
     );
-    clients.elasticAssistant.getCurrentUser.mockResolvedValue({
+    context.elasticAssistant.getCurrentUser.mockReturnValue({
       username: 'my_username',
       authentication_realm: {
         type: 'my_realm_type',
         name: 'my_realm_name',
       },
-    });
+    } as AuthenticatedUser);
 
     findUserConversationsRoute(server.router);
   });