diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/entries/use_create_knowledge_base_entry.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/entries/use_create_knowledge_base_entry.test.tsx new file mode 100644 index 0000000000000..73d0ddb976861 --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/entries/use_create_knowledge_base_entry.test.tsx @@ -0,0 +1,109 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { act, renderHook } from '@testing-library/react-hooks'; +import { + useCreateKnowledgeBaseEntry, + UseCreateKnowledgeBaseEntryParams, +} from './use_create_knowledge_base_entry'; +import { useInvalidateKnowledgeBaseEntries } from './use_knowledge_base_entries'; + +jest.mock('./use_knowledge_base_entries', () => ({ + useInvalidateKnowledgeBaseEntries: jest.fn(), +})); + +jest.mock('@tanstack/react-query', () => ({ + useMutation: jest.fn().mockImplementation((queryKey, fn, opts) => { + return { + mutate: async (variables: unknown) => { + try { + const res = await fn(variables); + opts.onSuccess(res); + opts.onSettled(); + return Promise.resolve(res); + } catch (e) { + opts.onError(e); + opts.onSettled(); + } + }, + }; + }), +})); + +const http = { + post: jest.fn(), +}; +const toasts = { + addError: jest.fn(), + addSuccess: jest.fn(), +}; +const defaultProps = { http, toasts } as unknown as UseCreateKnowledgeBaseEntryParams; +const defaultArgs = { title: 'Test Entry' }; +describe('useCreateKnowledgeBaseEntry', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('should call the mutation function on success', async () => { + const invalidateKnowledgeBaseEntries = jest.fn(); + (useInvalidateKnowledgeBaseEntries as jest.Mock).mockReturnValue( + invalidateKnowledgeBaseEntries + ); + http.post.mockResolvedValue({}); + + const { result } = renderHook(() => useCreateKnowledgeBaseEntry(defaultProps)); + + await act(async () => { + // @ts-ignore + await result.current.mutate(defaultArgs); + }); + + expect(http.post).toHaveBeenCalledWith( + expect.any(String), + expect.objectContaining({ + body: JSON.stringify(defaultArgs), + }) + ); + expect(toasts.addSuccess).toHaveBeenCalledWith({ + title: expect.any(String), + }); + expect(invalidateKnowledgeBaseEntries).toHaveBeenCalled(); + }); + + it('should call the onError function on error', async () => { + const error = new Error('Test Error'); + http.post.mockRejectedValue(error); + + const { result } = renderHook(() => useCreateKnowledgeBaseEntry(defaultProps)); + + await act(async () => { + // @ts-ignore + await result.current.mutate(defaultArgs); + }); + + expect(toasts.addError).toHaveBeenCalledWith(error, { + title: expect.any(String), + }); + }); + + it('should call the onSettled function after mutation', async () => { + const invalidateKnowledgeBaseEntries = jest.fn(); + (useInvalidateKnowledgeBaseEntries as jest.Mock).mockReturnValue( + invalidateKnowledgeBaseEntries + ); + http.post.mockResolvedValue({}); + + const { result } = renderHook(() => useCreateKnowledgeBaseEntry(defaultProps)); + + await act(async () => { + // @ts-ignore + await result.current.mutate(defaultArgs); + }); + + expect(invalidateKnowledgeBaseEntries).toHaveBeenCalled(); + }); +}); diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/entries/use_delete_knowledge_base_entries.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/entries/use_delete_knowledge_base_entries.test.tsx new file mode 100644 index 0000000000000..6003b1f81f435 --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/entries/use_delete_knowledge_base_entries.test.tsx @@ -0,0 +1,107 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { act, renderHook } from '@testing-library/react-hooks'; +import { + useDeleteKnowledgeBaseEntries, + UseDeleteKnowledgeEntriesParams, +} from './use_delete_knowledge_base_entries'; +import { useInvalidateKnowledgeBaseEntries } from './use_knowledge_base_entries'; + +jest.mock('./use_knowledge_base_entries', () => ({ + useInvalidateKnowledgeBaseEntries: jest.fn(), +})); + +jest.mock('@tanstack/react-query', () => ({ + useMutation: jest.fn().mockImplementation((queryKey, fn, opts) => { + return { + mutate: async (variables: unknown) => { + try { + const res = await fn(variables); + opts.onSuccess(res); + opts.onSettled(); + return Promise.resolve(res); + } catch (e) { + opts.onError(e); + opts.onSettled(); + } + }, + }; + }), +})); + +const http = { + post: jest.fn(), +}; +const toasts = { + addError: jest.fn(), +}; +const defaultProps = { http, toasts } as unknown as UseDeleteKnowledgeEntriesParams; +const defaultArgs = { ids: ['1'], query: '' }; + +describe('useDeleteKnowledgeBaseEntries', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('should call the mutation function on success', async () => { + const invalidateKnowledgeBaseEntries = jest.fn(); + (useInvalidateKnowledgeBaseEntries as jest.Mock).mockReturnValue( + invalidateKnowledgeBaseEntries + ); + http.post.mockResolvedValue({}); + + const { result } = renderHook(() => useDeleteKnowledgeBaseEntries(defaultProps)); + + await act(async () => { + // @ts-ignore + await result.current.mutate(defaultArgs); + }); + + expect(http.post).toHaveBeenCalledWith( + expect.any(String), + expect.objectContaining({ + body: JSON.stringify({ delete: { query: '', ids: ['1'] } }), + version: '1', + }) + ); + expect(invalidateKnowledgeBaseEntries).toHaveBeenCalled(); + }); + + it('should call the onError function on error', async () => { + const error = new Error('Test Error'); + http.post.mockRejectedValue(error); + + const { result } = renderHook(() => useDeleteKnowledgeBaseEntries(defaultProps)); + + await act(async () => { + // @ts-ignore + await result.current.mutate(defaultArgs); + }); + + expect(toasts.addError).toHaveBeenCalledWith(error, { + title: expect.any(String), + }); + }); + + it('should call the onSettled function after mutation', async () => { + const invalidateKnowledgeBaseEntries = jest.fn(); + (useInvalidateKnowledgeBaseEntries as jest.Mock).mockReturnValue( + invalidateKnowledgeBaseEntries + ); + http.post.mockResolvedValue({}); + + const { result } = renderHook(() => useDeleteKnowledgeBaseEntries(defaultProps)); + + await act(async () => { + // @ts-ignore + await result.current.mutate(defaultArgs); + }); + + expect(invalidateKnowledgeBaseEntries).toHaveBeenCalled(); + }); +}); diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/entries/use_knowledge_base_entries.test.ts b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/entries/use_knowledge_base_entries.test.ts new file mode 100644 index 0000000000000..f298258800d35 --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/entries/use_knowledge_base_entries.test.ts @@ -0,0 +1,76 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { renderHook } from '@testing-library/react-hooks'; +import { useKnowledgeBaseEntries } from './use_knowledge_base_entries'; +import { HttpSetup } from '@kbn/core/public'; +import { IToasts } from '@kbn/core-notifications-browser'; +import { TestProviders } from '../../../../mock/test_providers/test_providers'; + +describe('useKnowledgeBaseEntries', () => { + const httpMock: HttpSetup = { + fetch: jest.fn(), + } as unknown as HttpSetup; + const toastsMock: IToasts = { + addError: jest.fn(), + } as unknown as IToasts; + + it('fetches knowledge base entries successfully', async () => { + (httpMock.fetch as jest.Mock).mockResolvedValue({ + page: 1, + perPage: 100, + total: 1, + data: [{ id: '1', title: 'Entry 1' }], + }); + + const { result, waitForNextUpdate } = renderHook( + () => useKnowledgeBaseEntries({ http: httpMock, enabled: true }), + { + wrapper: TestProviders, + } + ); + expect(result.current.fetchStatus).toEqual('fetching'); + + await waitForNextUpdate(); + + expect(result.current.data).toEqual({ + page: 1, + perPage: 100, + total: 1, + data: [{ id: '1', title: 'Entry 1' }], + }); + }); + + it('handles fetch error', async () => { + const error = new Error('Fetch error'); + (httpMock.fetch as jest.Mock).mockRejectedValue(error); + + const { waitForNextUpdate } = renderHook( + () => useKnowledgeBaseEntries({ http: httpMock, toasts: toastsMock, enabled: true }), + { + wrapper: TestProviders, + } + ); + + await waitForNextUpdate(); + + expect(toastsMock.addError).toHaveBeenCalledWith(error, { + title: 'Error fetching Knowledge Base entries', + }); + }); + + it('does not fetch when disabled', async () => { + const { result } = renderHook( + () => useKnowledgeBaseEntries({ http: httpMock, enabled: false }), + { + wrapper: TestProviders, + } + ); + + expect(result.current.fetchStatus).toEqual('idle'); + }); +}); diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/entries/use_update_knowledge_base_entries.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/entries/use_update_knowledge_base_entries.test.tsx new file mode 100644 index 0000000000000..0c35727846a01 --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/entries/use_update_knowledge_base_entries.test.tsx @@ -0,0 +1,111 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { act, renderHook } from '@testing-library/react-hooks'; +import { + useUpdateKnowledgeBaseEntries, + UseUpdateKnowledgeBaseEntriesParams, +} from './use_update_knowledge_base_entries'; +import { useInvalidateKnowledgeBaseEntries } from './use_knowledge_base_entries'; + +jest.mock('./use_knowledge_base_entries', () => ({ + useInvalidateKnowledgeBaseEntries: jest.fn(), +})); + +jest.mock('@tanstack/react-query', () => ({ + useMutation: jest.fn().mockImplementation((queryKey, fn, opts) => { + return { + mutate: async (variables: unknown) => { + try { + const res = await fn(variables); + opts.onSuccess(res); + opts.onSettled(); + return Promise.resolve(res); + } catch (e) { + opts.onError(e); + opts.onSettled(); + } + }, + }; + }), +})); + +const http = { + post: jest.fn(), +}; +const toasts = { + addError: jest.fn(), + addSuccess: jest.fn(), +}; +const defaultProps = { http, toasts } as unknown as UseUpdateKnowledgeBaseEntriesParams; +const defaultArgs = { ids: ['1'], query: '', data: { field: 'value' } }; + +describe('useUpdateKnowledgeBaseEntries', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('should call the mutation function on success', async () => { + const invalidateKnowledgeBaseEntries = jest.fn(); + (useInvalidateKnowledgeBaseEntries as jest.Mock).mockReturnValue( + invalidateKnowledgeBaseEntries + ); + http.post.mockResolvedValue({}); + + const { result } = renderHook(() => useUpdateKnowledgeBaseEntries(defaultProps)); + + await act(async () => { + // @ts-ignore + await result.current.mutate(defaultArgs); + }); + + expect(http.post).toHaveBeenCalledWith( + expect.any(String), + expect.objectContaining({ + body: JSON.stringify({ update: defaultArgs }), + version: '1', + }) + ); + expect(invalidateKnowledgeBaseEntries).toHaveBeenCalled(); + expect(toasts.addSuccess).toHaveBeenCalledWith({ + title: expect.any(String), + }); + }); + + it('should call the onError function on error', async () => { + const error = new Error('Test Error'); + http.post.mockRejectedValue(error); + + const { result } = renderHook(() => useUpdateKnowledgeBaseEntries(defaultProps)); + + await act(async () => { + // @ts-ignore + await result.current.mutate(defaultArgs); + }); + + expect(toasts.addError).toHaveBeenCalledWith(error, { + title: expect.any(String), + }); + }); + + it('should call the onSettled function after mutation', async () => { + const invalidateKnowledgeBaseEntries = jest.fn(); + (useInvalidateKnowledgeBaseEntries as jest.Mock).mockReturnValue( + invalidateKnowledgeBaseEntries + ); + http.post.mockResolvedValue({}); + + const { result } = renderHook(() => useUpdateKnowledgeBaseEntries(defaultProps)); + + await act(async () => { + // @ts-ignore + await result.current.mutate(defaultArgs); + }); + + expect(invalidateKnowledgeBaseEntries).toHaveBeenCalled(); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/__mocks__/data_clients.mock.ts b/x-pack/plugins/elastic_assistant/server/__mocks__/data_clients.mock.ts index 473965a835f14..7c4abffff6520 100644 --- a/x-pack/plugins/elastic_assistant/server/__mocks__/data_clients.mock.ts +++ b/x-pack/plugins/elastic_assistant/server/__mocks__/data_clients.mock.ts @@ -7,6 +7,7 @@ import type { PublicMethodsOf } from '@kbn/utility-types'; import { AIAssistantConversationsDataClient } from '../ai_assistant_data_clients/conversations'; +import { AIAssistantKnowledgeBaseDataClient } from '../ai_assistant_data_clients/knowledge_base'; import { AIAssistantDataClient } from '../ai_assistant_data_clients'; import { AttackDiscoveryDataClient } from '../lib/attack_discovery/persistence'; @@ -14,6 +15,8 @@ type ConversationsDataClientContract = PublicMethodsOf; type AttackDiscoveryDataClientContract = PublicMethodsOf; export type AttackDiscoveryDataClientMock = jest.Mocked; +type KnowledgeBaseDataClientContract = PublicMethodsOf; +export type KnowledgeBaseDataClientMock = jest.Mocked; const createConversationsDataClientMock = () => { const mocked: ConversationsDataClientMock = { @@ -52,6 +55,33 @@ export const attackDiscoveryDataClientMock: { create: createAttackDiscoveryDataClientMock, }; +const createKnowledgeBaseDataClientMock = () => { + const mocked: KnowledgeBaseDataClientMock = { + addKnowledgeBaseDocuments: jest.fn(), + createInferenceEndpoint: jest.fn(), + createKnowledgeBaseEntry: jest.fn(), + findDocuments: jest.fn(), + getAssistantTools: jest.fn(), + getKnowledgeBaseDocumentEntries: jest.fn(), + getReader: jest.fn(), + getRequiredKnowledgeBaseDocumentEntries: jest.fn(), + getWriter: jest.fn().mockResolvedValue({ bulk: jest.fn() }), + isInferenceEndpointExists: jest.fn(), + isModelInstalled: jest.fn(), + isSecurityLabsDocsLoaded: jest.fn(), + isSetupAvailable: jest.fn(), + isUserDataExists: jest.fn(), + setupKnowledgeBase: jest.fn(), + }; + return mocked; +}; + +export const knowledgeBaseDataClientMock: { + create: () => KnowledgeBaseDataClientMock; +} = { + create: createKnowledgeBaseDataClientMock, +}; + type AIAssistantDataClientContract = PublicMethodsOf; export type AIAssistantDataClientMock = jest.Mocked; diff --git a/x-pack/plugins/elastic_assistant/server/__mocks__/knowledge_base_entry_schema.mock.ts b/x-pack/plugins/elastic_assistant/server/__mocks__/knowledge_base_entry_schema.mock.ts new file mode 100644 index 0000000000000..8171dd2b39249 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/__mocks__/knowledge_base_entry_schema.mock.ts @@ -0,0 +1,174 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { estypes } from '@elastic/elasticsearch'; +import { + KnowledgeBaseEntryCreateProps, + KnowledgeBaseEntryResponse, + KnowledgeBaseEntryUpdateProps, +} from '@kbn/elastic-assistant-common'; +import { + EsKnowledgeBaseEntrySchema, + EsDocumentEntry, +} from '../ai_assistant_data_clients/knowledge_base/types'; +const indexEntry: EsKnowledgeBaseEntrySchema = { + id: '1234', + '@timestamp': '2020-04-20T15:25:31.830Z', + created_at: '2020-04-20T15:25:31.830Z', + created_by: 'my_profile_uid', + updated_at: '2020-04-20T15:25:31.830Z', + updated_by: 'my_profile_uid', + name: 'test', + namespace: 'default', + type: 'index', + index: 'test', + field: 'test', + description: 'test', + query_description: 'test', + input_schema: [ + { + field_name: 'test', + field_type: 'test', + description: 'test', + }, + ], + users: [ + { + name: 'my_username', + id: 'my_profile_uid', + }, + ], +}; +export const documentEntry: EsDocumentEntry = { + id: '5678', + '@timestamp': '2020-04-20T15:25:31.830Z', + created_at: '2020-04-20T15:25:31.830Z', + created_by: 'my_profile_uid', + updated_at: '2020-04-20T15:25:31.830Z', + updated_by: 'my_profile_uid', + name: 'test', + namespace: 'default', + semantic_text: 'test', + type: 'document', + kb_resource: 'test', + required: true, + source: 'test', + text: 'test', + users: [ + { + name: 'my_username', + id: 'my_profile_uid', + }, + ], +}; + +export const getKnowledgeBaseEntrySearchEsMock = (src = 'document') => { + const searchResponse: estypes.SearchResponse = { + took: 3, + timed_out: false, + _shards: { + total: 2, + successful: 2, + skipped: 0, + failed: 0, + }, + hits: { + total: { + value: 1, + relation: 'eq', + }, + max_score: 0, + hits: [ + { + _id: '1', + _index: '', + _score: 0, + _source: src === 'document' ? documentEntry : indexEntry, + }, + ], + }, + }; + return searchResponse; +}; + +export const getCreateKnowledgeBaseEntrySchemaMock = ( + rest?: Partial +): KnowledgeBaseEntryCreateProps => { + const { type = 'document', ...restProps } = rest ?? {}; + if (type === 'document') { + return { + type: 'document', + source: 'test', + text: 'test', + name: 'test', + kbResource: 'test', + ...restProps, + }; + } + return { + type: 'index', + name: 'test', + index: 'test', + field: 'test', + description: 'test', + queryDescription: 'test', + inputSchema: [ + { + fieldName: 'test', + fieldType: 'test', + description: 'test', + }, + ], + ...restProps, + }; +}; + +export const getUpdateKnowledgeBaseEntrySchemaMock = ( + entryId = 'entry-1' +): KnowledgeBaseEntryUpdateProps => ({ + name: 'another 2', + namespace: 'default', + type: 'document', + source: 'test', + text: 'test', + kbResource: 'test', + id: entryId, +}); + +export const getKnowledgeBaseEntryMock = ( + params: KnowledgeBaseEntryCreateProps | KnowledgeBaseEntryUpdateProps = { + name: 'test', + namespace: 'default', + type: 'document', + text: 'test', + source: 'test', + kbResource: 'test', + required: true, + } +): KnowledgeBaseEntryResponse => ({ + id: '1', + ...params, + createdBy: 'my_profile_uid', + updatedBy: 'my_profile_uid', + createdAt: '2020-04-20T15:25:31.830Z', + updatedAt: '2020-04-20T15:25:31.830Z', + namespace: 'default', + users: [ + { + name: 'my_username', + id: 'my_profile_uid', + }, + ], +}); + +export const getQueryKnowledgeBaseEntryParams = ( + isUpdate?: boolean +): KnowledgeBaseEntryCreateProps | KnowledgeBaseEntryUpdateProps => { + return isUpdate + ? getUpdateKnowledgeBaseEntrySchemaMock() + : getCreateKnowledgeBaseEntrySchemaMock(); +}; diff --git a/x-pack/plugins/elastic_assistant/server/__mocks__/request.ts b/x-pack/plugins/elastic_assistant/server/__mocks__/request.ts index 698645e8d3c55..b62cd24e938eb 100644 --- a/x-pack/plugins/elastic_assistant/server/__mocks__/request.ts +++ b/x-pack/plugins/elastic_assistant/server/__mocks__/request.ts @@ -23,10 +23,14 @@ import { ELASTIC_AI_ASSISTANT_CONVERSATIONS_URL_BY_ID_MESSAGES, ELASTIC_AI_ASSISTANT_CONVERSATIONS_URL_FIND, ELASTIC_AI_ASSISTANT_EVALUATE_URL, + ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL, + ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_BULK_ACTION, + ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_FIND, ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_INDICES_URL, ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL, ELASTIC_AI_ASSISTANT_PROMPTS_URL_BULK_ACTION, ELASTIC_AI_ASSISTANT_PROMPTS_URL_FIND, + PerformKnowledgeBaseEntryBulkActionRequestBody, PostEvaluateRequestBodyInput, } from '@kbn/elastic-assistant-common'; import { @@ -34,6 +38,7 @@ import { getCreateConversationSchemaMock, getUpdateConversationSchemaMock, } from './conversations_schema.mock'; +import { getCreateKnowledgeBaseEntrySchemaMock } from './knowledge_base_entry_schema.mock'; import { PromptCreateProps, PromptUpdateProps, @@ -67,6 +72,22 @@ export const getPostKnowledgeBaseRequest = (resource?: string) => query: { resource }, }); +export const getCreateKnowledgeBaseEntryRequest = () => + requestMock.create({ + method: 'post', + path: ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL, + body: getCreateKnowledgeBaseEntrySchemaMock(), + }); + +export const getBulkActionKnowledgeBaseEntryRequest = ( + body: PerformKnowledgeBaseEntryBulkActionRequestBody +) => + requestMock.create({ + method: 'post', + path: ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_BULK_ACTION, + body, + }); + export const getGetCapabilitiesRequest = () => requestMock.create({ method: 'get', @@ -80,6 +101,12 @@ export const getPostEvaluateRequest = ({ body }: { body: PostEvaluateRequestBody path: ELASTIC_AI_ASSISTANT_EVALUATE_URL, }); +export const getKnowledgeBaseEntryFindRequest = () => + requestMock.create({ + method: 'get', + path: ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_FIND, + }); + export const getCurrentUserFindRequest = () => requestMock.create({ method: 'get', diff --git a/x-pack/plugins/elastic_assistant/server/__mocks__/request_context.ts b/x-pack/plugins/elastic_assistant/server/__mocks__/request_context.ts index a065c7de42586..19d98633a83c9 100644 --- a/x-pack/plugins/elastic_assistant/server/__mocks__/request_context.ts +++ b/x-pack/plugins/elastic_assistant/server/__mocks__/request_context.ts @@ -18,6 +18,7 @@ import { attackDiscoveryDataClientMock, conversationsDataClientMock, dataClientMock, + knowledgeBaseDataClientMock, } from './data_clients.mock'; import { AIAssistantConversationsDataClient } from '../ai_assistant_data_clients/conversations'; import { AIAssistantDataClient } from '../ai_assistant_data_clients'; @@ -27,6 +28,7 @@ import { } from '../ai_assistant_data_clients/knowledge_base'; import { defaultAssistantFeatures } from '@kbn/elastic-assistant-common'; import { AttackDiscoveryDataClient } from '../lib/attack_discovery/persistence'; +import { authenticatedUser } from './user'; export const createMockClients = () => { const core = coreMock.createRequestHandlerContext(); @@ -42,7 +44,7 @@ export const createMockClients = () => { logger: loggingSystemMock.createLogger(), telemetry: coreMock.createSetup().analytics, getAIAssistantConversationsDataClient: conversationsDataClientMock.create(), - getAIAssistantKnowledgeBaseDataClient: dataClientMock.create(), + getAIAssistantKnowledgeBaseDataClient: knowledgeBaseDataClientMock.create(), getAIAssistantPromptsDataClient: dataClientMock.create(), getAttackDiscoveryDataClient: attackDiscoveryDataClientMock.create(), getAIAssistantAnonymizationFieldsDataClient: dataClientMock.create(), @@ -133,9 +135,9 @@ const createElasticAssistantRequestContextMock = ( (( params?: GetAIAssistantKnowledgeBaseDataClientParams ) => Promise), - getCurrentUser: jest.fn(), + getCurrentUser: jest.fn().mockReturnValue(authenticatedUser), getServerBasePath: jest.fn(), - getSpaceId: jest.fn(), + getSpaceId: jest.fn().mockReturnValue('default'), inference: { getClient: jest.fn() }, core: clients.core, telemetry: clients.elasticAssistant.telemetry, diff --git a/x-pack/plugins/elastic_assistant/server/__mocks__/response.ts b/x-pack/plugins/elastic_assistant/server/__mocks__/response.ts index dc5a2ba0e884a..b7ab289d0f270 100644 --- a/x-pack/plugins/elastic_assistant/server/__mocks__/response.ts +++ b/x-pack/plugins/elastic_assistant/server/__mocks__/response.ts @@ -15,6 +15,8 @@ import { EsPromptsSchema } from '../ai_assistant_data_clients/prompts/types'; import { getPromptsSearchEsMock } from './prompts_schema.mock'; import { EsAnonymizationFieldsSchema } from '../ai_assistant_data_clients/anonymization_fields/types'; import { getAnonymizationFieldsSearchEsMock } from './anonymization_fields_schema.mock'; +import { getKnowledgeBaseEntrySearchEsMock } from './knowledge_base_entry_schema.mock'; +import { EsKnowledgeBaseEntrySchema } from '../ai_assistant_data_clients/knowledge_base/types'; export const responseMock = { create: httpServerMock.createResponseFactory, @@ -27,6 +29,14 @@ export const getEmptyFindResult = (): FindResponse => ({ data: getBasicEmptySearchResponse(), }); +export const getFindKnowledgeBaseEntriesResultWithSingleHit = + (): FindResponse => ({ + page: 1, + perPage: 1, + total: 1, + data: getKnowledgeBaseEntrySearchEsMock(), + }); + export const getFindConversationsResultWithSingleHit = (): FindResponse => ({ page: 1, perPage: 1, diff --git a/x-pack/plugins/elastic_assistant/server/__mocks__/user.ts b/x-pack/plugins/elastic_assistant/server/__mocks__/user.ts new file mode 100644 index 0000000000000..bcd29818c4ed7 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/__mocks__/user.ts @@ -0,0 +1,17 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { AuthenticatedUser } from '@kbn/core-security-common'; + +export const authenticatedUser = { + username: 'my_username', + profile_uid: 'my_profile_uid', + authentication_realm: { + type: 'my_realm_type', + name: 'my_realm_name', + }, +} as AuthenticatedUser; diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/create_conversation.test.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/create_conversation.test.ts index 7ef1f7865da36..0546ab39db592 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/create_conversation.test.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/create_conversation.test.ts @@ -9,20 +9,14 @@ import { elasticsearchClientMock } from '@kbn/core-elasticsearch-client-server-m import { createConversation } from './create_conversation'; import { loggingSystemMock } from '@kbn/core-logging-server-mocks'; import { getConversation } from './get_conversation'; +import { authenticatedUser } from '../../__mocks__/user'; import { ConversationCreateProps, ConversationResponse } from '@kbn/elastic-assistant-common'; -import { AuthenticatedUser } from '@kbn/core-security-common'; jest.mock('./get_conversation', () => ({ getConversation: jest.fn(), })); -const mockUser1 = { - username: 'my_username', - authentication_realm: { - type: 'my_realm_type', - name: 'my_realm_name', - }, -} as AuthenticatedUser; +const mockUser1 = authenticatedUser; export const getCreateConversationMock = (): ConversationCreateProps => ({ title: 'test', diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/get_conversation.test.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/get_conversation.test.ts index a5a292c096cdc..43290c8a00293 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/get_conversation.test.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/get_conversation.test.ts @@ -5,11 +5,12 @@ * 2.0. */ -import type { AuthenticatedUser, Logger } from '@kbn/core/server'; +import type { Logger } from '@kbn/core/server'; import { elasticsearchClientMock } from '@kbn/core-elasticsearch-client-server-mocks'; import { getConversation } from './get_conversation'; import { estypes } from '@elastic/elasticsearch'; import { EsConversationSchema } from './types'; +import { authenticatedUser } from '../../__mocks__/user'; import { loggingSystemMock } from '@kbn/core-logging-server-mocks'; import { ConversationResponse } from '@kbn/elastic-assistant-common'; @@ -43,13 +44,7 @@ export const getConversationResponseMock = (): ConversationResponse => ({ replacements: undefined, }); -const mockUser1 = { - username: 'my_username', - authentication_realm: { - type: 'my_realm_type', - name: 'my_realm_name', - }, -} as AuthenticatedUser; +const mockUser1 = authenticatedUser; export const getSearchConversationMock = (): estypes.SearchResponse => ({ _scroll_id: '123', diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/index.test.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/index.test.ts index 7669c281a42da..4c57f66710f5e 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/index.test.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/index.test.ts @@ -7,21 +7,15 @@ import { elasticsearchServiceMock, loggingSystemMock } from '@kbn/core/server/mocks'; import type { UpdateByQueryRequest } from '@elastic/elasticsearch/lib/api/types'; import { AIAssistantConversationsDataClient } from '.'; -import { AuthenticatedUser } from '@kbn/core-security-common'; import { getUpdateConversationSchemaMock } from '../../__mocks__/conversations_schema.mock'; +import { authenticatedUser } from '../../__mocks__/user'; import { AIAssistantDataClientParams } from '..'; const date = '2023-03-28T22:27:28.159Z'; let logger: ReturnType<(typeof loggingSystemMock)['createLogger']>; const clusterClient = elasticsearchServiceMock.createClusterClient().asInternalUser; -const mockUser1 = { - username: 'my_username', - authentication_realm: { - type: 'my_realm_type', - name: 'my_realm_name', - }, -} as AuthenticatedUser; +const mockUser1 = authenticatedUser; describe('AIAssistantConversationsDataClient', () => { let assistantConversationsDataClientParams: AIAssistantDataClientParams; diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/update_conversation.test.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/update_conversation.test.ts index c44329c28db48..baeea677b1a66 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/update_conversation.test.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/conversations/update_conversation.test.ts @@ -13,8 +13,8 @@ import { updateConversation, } from './update_conversation'; import { getConversation } from './get_conversation'; +import { authenticatedUser } from '../../__mocks__/user'; import { ConversationResponse, ConversationUpdateProps } from '@kbn/elastic-assistant-common'; -import { AuthenticatedUser } from '@kbn/core-security-common'; export const getUpdateConversationOptionsMock = (): ConversationUpdateProps => ({ id: 'test', @@ -31,13 +31,7 @@ export const getUpdateConversationOptionsMock = (): ConversationUpdateProps => ( replacements: {}, }); -const mockUser1 = { - username: 'my_username', - authentication_realm: { - type: 'my_realm_type', - name: 'my_realm_name', - }, -} as AuthenticatedUser; +const mockUser1 = authenticatedUser; export const getConversationResponseMock = (): ConversationResponse => ({ id: 'test', diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/index.test.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/index.test.ts index 8167708431921..007e25e9af467 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/index.test.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/index.test.ts @@ -6,19 +6,12 @@ */ import { elasticsearchServiceMock, loggingSystemMock } from '@kbn/core/server/mocks'; import { AIAssistantDataClient, AIAssistantDataClientParams } from '.'; -import { AuthenticatedUser } from '@kbn/core-security-common'; - +import { authenticatedUser } from '../__mocks__/user'; const date = '2023-03-28T22:27:28.159Z'; let logger: ReturnType<(typeof loggingSystemMock)['createLogger']>; const clusterClient = elasticsearchServiceMock.createClusterClient().asInternalUser; -const mockUser1 = { - username: 'my_username', - authentication_realm: { - type: 'my_realm_type', - name: 'my_realm_name', - }, -} as AuthenticatedUser; +const mockUser1 = authenticatedUser; describe('AIAssistantDataClient', () => { let assistantDataClientParams: AIAssistantDataClientParams; diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/create_knowledge_base_entry.test.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/create_knowledge_base_entry.test.ts new file mode 100644 index 0000000000000..df6533d5d8df2 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/create_knowledge_base_entry.test.ts @@ -0,0 +1,172 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { elasticsearchClientMock } from '@kbn/core-elasticsearch-client-server-mocks'; +import { createKnowledgeBaseEntry } from './create_knowledge_base_entry'; +import { loggingSystemMock } from '@kbn/core-logging-server-mocks'; +import { coreMock } from '@kbn/core/server/mocks'; +import { getKnowledgeBaseEntry } from './get_knowledge_base_entry'; +import { KnowledgeBaseEntryResponse } from '@kbn/elastic-assistant-common'; +import { + getKnowledgeBaseEntryMock, + getCreateKnowledgeBaseEntrySchemaMock, +} from '../../__mocks__/knowledge_base_entry_schema.mock'; +import { authenticatedUser } from '../../__mocks__/user'; + +jest.mock('./get_knowledge_base_entry', () => ({ + getKnowledgeBaseEntry: jest.fn(), +})); + +const telemetry = coreMock.createSetup().analytics; + +describe('createKnowledgeBaseEntry', () => { + let logger: ReturnType; + beforeEach(() => { + jest.clearAllMocks(); + logger = loggingSystemMock.createLogger(); + }); + + afterEach(() => { + jest.clearAllMocks(); + }); + + beforeAll(() => { + jest.useFakeTimers(); + const date = '2024-01-28T04:20:02.394Z'; + jest.setSystemTime(new Date(date)); + }); + + afterAll(() => { + jest.useRealTimers(); + }); + + test('it creates a knowledge base document entry with create schema', async () => { + const knowledgeBaseEntry = getCreateKnowledgeBaseEntrySchemaMock(); + (getKnowledgeBaseEntry as unknown as jest.Mock).mockResolvedValueOnce({ + ...getKnowledgeBaseEntryMock(), + id: 'elastic-id-123', + }); + + const esClient = elasticsearchClientMock.createScopedClusterClient().asCurrentUser; + esClient.create.mockResponse( + // @ts-expect-error not full response interface + { _id: 'elastic-id-123' } + ); + const createdEntry = await createKnowledgeBaseEntry({ + esClient, + knowledgeBaseIndex: 'index-1', + spaceId: 'test', + user: authenticatedUser, + knowledgeBaseEntry, + logger, + telemetry, + }); + expect(esClient.create).toHaveBeenCalledWith({ + body: { + '@timestamp': '2024-01-28T04:20:02.394Z', + created_at: '2024-01-28T04:20:02.394Z', + created_by: 'my_profile_uid', + updated_at: '2024-01-28T04:20:02.394Z', + updated_by: 'my_profile_uid', + namespace: 'test', + users: [{ id: 'my_profile_uid', name: 'my_username' }], + type: 'document', + semantic_text: 'test', + source: 'test', + text: 'test', + name: 'test', + kb_resource: 'test', + required: false, + vector: undefined, + }, + id: expect.any(String), + index: 'index-1', + refresh: 'wait_for', + }); + + const expected: KnowledgeBaseEntryResponse = { + ...getKnowledgeBaseEntryMock(), + id: 'elastic-id-123', + }; + + expect(createdEntry).toEqual(expected); + }); + + test('it creates a knowledge base index entry with create schema', async () => { + const knowledgeBaseEntry = getCreateKnowledgeBaseEntrySchemaMock({ type: 'index' }); + (getKnowledgeBaseEntry as unknown as jest.Mock).mockResolvedValueOnce({ + ...getKnowledgeBaseEntryMock(), + id: 'elastic-id-123', + }); + + const esClient = elasticsearchClientMock.createScopedClusterClient().asCurrentUser; + esClient.create.mockResponse( + // @ts-expect-error not full response interface + { _id: 'elastic-id-123' } + ); + const createdEntry = await createKnowledgeBaseEntry({ + esClient, + knowledgeBaseIndex: 'index-1', + spaceId: 'test', + user: authenticatedUser, + knowledgeBaseEntry, + logger, + telemetry, + }); + expect(esClient.create).toHaveBeenCalledWith({ + body: { + '@timestamp': '2024-01-28T04:20:02.394Z', + created_at: '2024-01-28T04:20:02.394Z', + created_by: 'my_profile_uid', + updated_at: '2024-01-28T04:20:02.394Z', + updated_by: 'my_profile_uid', + namespace: 'test', + users: [{ id: 'my_profile_uid', name: 'my_username' }], + query_description: 'test', + type: 'index', + name: 'test', + description: 'test', + field: 'test', + index: 'test', + input_schema: [ + { + description: 'test', + field_name: 'test', + field_type: 'test', + }, + ], + }, + id: expect.any(String), + index: 'index-1', + refresh: 'wait_for', + }); + + const expected: KnowledgeBaseEntryResponse = { + ...getKnowledgeBaseEntryMock(), + id: 'elastic-id-123', + }; + + expect(createdEntry).toEqual(expected); + }); + + test('it throws an error when creating a knowledge base entry fails', async () => { + const knowledgeBaseEntry = getCreateKnowledgeBaseEntrySchemaMock(); + const esClient = elasticsearchClientMock.createScopedClusterClient().asCurrentUser; + esClient.create.mockRejectedValue(new Error('Test error')); + await expect( + createKnowledgeBaseEntry({ + esClient, + knowledgeBaseIndex: 'index-1', + spaceId: 'test', + user: authenticatedUser, + knowledgeBaseEntry, + logger, + telemetry, + }) + ).rejects.toThrowError('Test error'); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/get_knowledge_base_entry.test.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/get_knowledge_base_entry.test.ts new file mode 100644 index 0000000000000..aa247ece22e9a --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/get_knowledge_base_entry.test.ts @@ -0,0 +1,74 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import type { AuthenticatedUser, Logger } from '@kbn/core/server'; +import { elasticsearchClientMock } from '@kbn/core-elasticsearch-client-server-mocks'; +import { getKnowledgeBaseEntry } from './get_knowledge_base_entry'; +import { loggingSystemMock } from '@kbn/core-logging-server-mocks'; + +import { + getKnowledgeBaseEntryMock, + getKnowledgeBaseEntrySearchEsMock, +} from '../../__mocks__/knowledge_base_entry_schema.mock'; +export const mockUser = { + username: 'my_username', + authentication_realm: { + type: 'my_realm_type', + name: 'my_realm_name', + }, +} as AuthenticatedUser; +describe('getKnowledgeBaseEntry', () => { + let loggerMock: Logger; + beforeEach(() => { + jest.clearAllMocks(); + loggerMock = loggingSystemMock.createLogger(); + }); + + test('it returns an entry as expected if the entry is found', async () => { + const data = getKnowledgeBaseEntrySearchEsMock(); + const esClient = elasticsearchClientMock.createScopedClusterClient().asCurrentUser; + esClient.search.mockResponse(data); + const entry = await getKnowledgeBaseEntry({ + esClient, + knowledgeBaseIndex: '.kibana-elastic-ai-assistant-knowledge-base', + id: '1', + logger: loggerMock, + user: mockUser, + }); + const expected = getKnowledgeBaseEntryMock(); + expect(entry).toEqual(expected); + }); + + test('it returns null if the search is empty', async () => { + const data = getKnowledgeBaseEntrySearchEsMock(); + data.hits.hits = []; + const esClient = elasticsearchClientMock.createScopedClusterClient().asCurrentUser; + esClient.search.mockResponse(data); + const entry = await getKnowledgeBaseEntry({ + esClient, + knowledgeBaseIndex: '.kibana-elastic-ai-assistant-knowledge-base', + id: '1', + logger: loggerMock, + user: mockUser, + }); + expect(entry).toEqual(null); + }); + + test('it throws an error if the search fails', async () => { + const esClient = elasticsearchClientMock.createScopedClusterClient().asCurrentUser; + esClient.search.mockRejectedValue(new Error('search failed')); + await expect( + getKnowledgeBaseEntry({ + esClient, + knowledgeBaseIndex: '.kibana-elastic-ai-assistant-knowledge-base', + id: '1', + logger: loggerMock, + user: mockUser, + }) + ).rejects.toThrowError('search failed'); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.test.tsx b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.test.tsx new file mode 100644 index 0000000000000..69b142bdac6be --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.test.tsx @@ -0,0 +1,234 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { errors } from '@elastic/elasticsearch'; +import { ElasticsearchClient, Logger } from '@kbn/core/server'; +import { DynamicStructuredTool } from '@langchain/core/tools'; +import { + isModelAlreadyExistsError, + getKBVectorSearchQuery, + getStructuredToolForIndexEntry, +} from './helpers'; +import { authenticatedUser } from '../../__mocks__/user'; +import { getCreateKnowledgeBaseEntrySchemaMock } from '../../__mocks__/knowledge_base_entry_schema.mock'; +import { IndexEntry } from '@kbn/elastic-assistant-common'; + +// Mock dependencies +jest.mock('@elastic/elasticsearch'); +jest.mock('@kbn/zod', () => ({ + z: { + string: jest.fn().mockReturnValue({ describe: (str: string) => str }), + number: jest.fn().mockReturnValue({ describe: (str: string) => str }), + boolean: jest.fn().mockReturnValue({ describe: (str: string) => str }), + object: jest.fn().mockReturnValue({ describe: (str: string) => str }), + any: jest.fn().mockReturnValue({ describe: (str: string) => str }), + }, +})); +jest.mock('lodash'); + +describe('isModelAlreadyExistsError', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + it('should return true if error is resource_not_found_exception', () => { + const error = new errors.ResponseError({ + meta: { + name: 'error', + context: 'error', + request: { + params: { method: 'post', path: '/' }, + options: {}, + id: 'error', + }, + connection: null, + attempts: 0, + aborted: false, + }, + warnings: null, + body: { error: { type: 'resource_not_found_exception' } }, + }); + // @ts-ignore + error.body = { + error: { + type: 'resource_not_found_exception', + }, + }; + expect(isModelAlreadyExistsError(error)).toBe(true); + }); + + it('should return true if error is status_exception', () => { + const error = new errors.ResponseError({ + meta: { + name: 'error', + context: 'error', + request: { + params: { method: 'post', path: '/' }, + options: {}, + id: 'error', + }, + connection: null, + attempts: 0, + aborted: false, + }, + warnings: null, + body: { error: { type: 'status_exception' } }, + }); + // @ts-ignore + error.body = { + error: { + type: 'status_exception', + }, + }; + expect(isModelAlreadyExistsError(error)).toBe(true); + }); + + it('should return false for other error types', () => { + const error = new Error('Some other error'); + expect(isModelAlreadyExistsError(error)).toBe(false); + }); +}); + +describe('getKBVectorSearchQuery', () => { + const mockUser = authenticatedUser; + + it('should construct a query with no filters if none are provided', () => { + const query = getKBVectorSearchQuery({ user: mockUser }); + expect(query).toEqual({ + bool: { + must: [], + should: expect.any(Array), + filter: undefined, + minimum_should_match: 1, + }, + }); + }); + + it('should include kbResource in the query if provided', () => { + const query = getKBVectorSearchQuery({ user: mockUser, kbResource: 'esql' }); + expect(query?.bool?.must).toEqual( + expect.arrayContaining([ + { + term: { kb_resource: 'esql' }, + }, + ]) + ); + }); + + it('should include required filter in the query if required is true', () => { + const query = getKBVectorSearchQuery({ user: mockUser, required: true }); + expect(query?.bool?.must).toEqual( + expect.arrayContaining([ + { + term: { required: true }, + }, + ]) + ); + }); + + it('should add semantic text filter if query is provided', () => { + const query = getKBVectorSearchQuery({ user: mockUser, query: 'example' }); + expect(query?.bool?.must).toEqual( + expect.arrayContaining([ + { + semantic: { + field: 'semantic_text', + query: 'example', + }, + }, + ]) + ); + }); +}); + +describe('getStructuredToolForIndexEntry', () => { + const mockLogger = { + debug: jest.fn(), + error: jest.fn(), + } as unknown as Logger; + + const mockEsClient = {} as ElasticsearchClient; + + const mockIndexEntry = getCreateKnowledgeBaseEntrySchemaMock({ type: 'index' }) as IndexEntry; + + it('should return a DynamicStructuredTool with correct name and schema', () => { + const tool = getStructuredToolForIndexEntry({ + indexEntry: mockIndexEntry, + esClient: mockEsClient, + logger: mockLogger, + elserId: 'elser123', + }); + + expect(tool).toBeInstanceOf(DynamicStructuredTool); + expect(tool.lc_kwargs).toEqual( + expect.objectContaining({ + name: 'test', + description: 'test', + tags: ['knowledge-base'], + }) + ); + }); + + it('should execute func correctly and return expected results', async () => { + const mockSearchResult = { + hits: { + hits: [ + { + _source: { + field1: 'value1', + field2: 2, + }, + inner_hits: { + 'test.test': { + hits: { + hits: [ + { _source: { text: 'Inner text 1' } }, + { _source: { text: 'Inner text 2' } }, + ], + }, + }, + }, + }, + ], + }, + }; + + mockEsClient.search = jest.fn().mockResolvedValue(mockSearchResult); + + const tool = getStructuredToolForIndexEntry({ + indexEntry: mockIndexEntry, + esClient: mockEsClient, + logger: mockLogger, + elserId: 'elser123', + }); + + const input = { query: 'testQuery', field1: 'value1', field2: 2 }; + const result = await tool.invoke(input, {}); + + expect(result).toContain('Below are all relevant documents in JSON format'); + expect(result).toContain('"text":"Inner text 1\\n --- \\nInner text 2"'); + }); + + it('should log an error and return error message on Elasticsearch error', async () => { + const mockError = new Error('Elasticsearch error'); + mockEsClient.search = jest.fn().mockRejectedValue(mockError); + + const tool = getStructuredToolForIndexEntry({ + indexEntry: mockIndexEntry, + esClient: mockEsClient, + logger: mockLogger, + elserId: 'elser123', + }); + + const input = { query: 'testQuery', field1: 'value1', field2: 2 }; + const result = await tool.invoke(input, {}); + + expect(mockLogger.error).toHaveBeenCalledWith( + `Error performing IndexEntry KB Similarity Search: ${mockError.message}` + ); + expect(result).toContain(`I'm sorry, but I was unable to find any information`); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts index 88ecae26cf19f..a0d3afb355b4b 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts @@ -173,13 +173,11 @@ export const getStructuredToolForIndexEntry = ({ // Generate filters for inputSchema fields const filter = - indexEntry.inputSchema?.reduce((prev, i) => { - return [ - ...prev, - // @ts-expect-error Possible to override types with dynamic input schema? - { term: { [`${i.fieldName}`]: input?.[i.fieldName] } }, - ]; - }, [] as Array<{ term: { [key: string]: string } }>) ?? []; + indexEntry.inputSchema?.reduce( + // @ts-expect-error Possible to override types with dynamic input schema? + (prev, i) => [...prev, { term: { [`${i.fieldName}`]: input?.[i.fieldName] } }], + [] as Array<{ term: { [key: string]: string } }> + ) ?? []; const params: SearchRequest = { index: indexEntry.index, diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.test.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.test.ts new file mode 100644 index 0000000000000..cf67d763e3d23 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.test.ts @@ -0,0 +1,582 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ +import { + coreMock, + elasticsearchServiceMock, + loggingSystemMock, + savedObjectsRepositoryMock, +} from '@kbn/core/server/mocks'; +import { AIAssistantKnowledgeBaseDataClient, KnowledgeBaseDataClientParams } from '.'; +import { + getCreateKnowledgeBaseEntrySchemaMock, + getKnowledgeBaseEntryMock, + getKnowledgeBaseEntrySearchEsMock, +} from '../../__mocks__/knowledge_base_entry_schema.mock'; +import { authenticatedUser } from '../../__mocks__/user'; +import { IndexPatternsFetcher } from '@kbn/data-plugin/server'; +import type { MlPluginSetup } from '@kbn/ml-plugin/server'; +import { mlPluginMock } from '@kbn/ml-plugin/public/mocks'; +import pRetry from 'p-retry'; + +import { + loadSecurityLabs, + getSecurityLabsDocsCount, +} from '../../lib/langchain/content_loaders/security_labs_loader'; +import { DynamicStructuredTool } from '@langchain/core/tools'; +jest.mock('../../lib/langchain/content_loaders/security_labs_loader'); +jest.mock('p-retry'); +const date = '2023-03-28T22:27:28.159Z'; +let logger: ReturnType<(typeof loggingSystemMock)['createLogger']>; +const esClientMock = elasticsearchServiceMock.createClusterClient().asInternalUser; + +const mockUser1 = authenticatedUser; + +const mockedPRetry = pRetry as jest.MockedFunction; +mockedPRetry.mockResolvedValue({}); +const telemetry = coreMock.createSetup().analytics; + +describe('AIAssistantKnowledgeBaseDataClient', () => { + let mockOptions: KnowledgeBaseDataClientParams; + let ml: MlPluginSetup; + let savedObjectClient: ReturnType; + const getElserId = jest.fn(); + const trainedModelsProvider = jest.fn(); + const installElasticModel = jest.fn(); + const mockLoadSecurityLabs = loadSecurityLabs as jest.Mock; + const mockGetSecurityLabsDocsCount = getSecurityLabsDocsCount as jest.Mock; + const mockGetIsKBSetupInProgress = jest.fn(); + beforeEach(() => { + jest.clearAllMocks(); + logger = loggingSystemMock.createLogger(); + savedObjectClient = savedObjectsRepositoryMock.create(); + mockLoadSecurityLabs.mockClear(); + ml = mlPluginMock.createSetupContract() as unknown as MlPluginSetup; // Missing SharedServices mock, so manually mocking trainedModelsProvider + ml.trainedModelsProvider = trainedModelsProvider.mockImplementation(() => ({ + getELSER: jest.fn().mockImplementation(() => '.elser_model_2'), + installElasticModel: installElasticModel.mockResolvedValue({}), + })); + mockOptions = { + logger, + elasticsearchClientPromise: Promise.resolve(esClientMock), + spaceId: 'default', + indexPatternsResourceName: '', + currentUser: mockUser1, + kibanaVersion: '8.8.0', + ml, + getElserId: getElserId.mockResolvedValue('elser-id'), + getIsKBSetupInProgress: mockGetIsKBSetupInProgress.mockReturnValue(false), + ingestPipelineResourceName: 'something', + setIsKBSetupInProgress: jest.fn().mockImplementation(() => {}), + manageGlobalKnowledgeBaseAIAssistant: true, + }; + esClientMock.search.mockReturnValue( + // @ts-expect-error not full response interface + getKnowledgeBaseEntrySearchEsMock() + ); + }); + + beforeAll(() => { + jest.useFakeTimers(); + jest.setSystemTime(new Date(date)); + }); + + afterAll(() => { + jest.useRealTimers(); + }); + describe('isSetupInProgress', () => { + it('should return true if setup is in progress', () => { + mockGetIsKBSetupInProgress.mockReturnValueOnce(true); + const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); + + const result = client.isSetupInProgress; + + expect(result).toBe(true); + }); + + it('should return false if setup is not in progress', () => { + mockGetIsKBSetupInProgress.mockReturnValueOnce(false); + const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); + + const result = client.isSetupInProgress; + + expect(result).toBe(false); + }); + }); + describe('isSetupAvailable', () => { + it('should return true if ML capabilities check succeeds', async () => { + const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); + // @ts-expect-error not full response interface + esClientMock.ml.getMemoryStats.mockResolvedValue({}); + const result = await client.isSetupAvailable(); + expect(result).toBe(true); + expect(esClientMock.ml.getMemoryStats).toHaveBeenCalled(); + }); + + it('should return false if ML capabilities check fails', async () => { + const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); + esClientMock.ml.getMemoryStats.mockRejectedValue(new Error('Mocked Error')); + const result = await client.isSetupAvailable(); + expect(result).toBe(false); + }); + }); + + describe('isModelInstalled', () => { + it('should check if ELSER model is installed and return true if fully_defined', async () => { + const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); + esClientMock.ml.getTrainedModels.mockResolvedValue({ + count: 1, + trained_model_configs: [ + { fully_defined: true, model_id: '', tags: [], input: { field_names: ['content'] } }, + ], + }); + const result = await client.isModelInstalled(); + expect(result).toBe(true); + expect(esClientMock.ml.getTrainedModels).toHaveBeenCalledWith({ + model_id: 'elser-id', + include: 'definition_status', + }); + }); + + it('should return false if model is not fully defined', async () => { + const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); + esClientMock.ml.getTrainedModels.mockResolvedValue({ + count: 0, + trained_model_configs: [ + { fully_defined: false, model_id: '', tags: [], input: { field_names: ['content'] } }, + ], + }); + const result = await client.isModelInstalled(); + expect(result).toBe(false); + }); + + it('should return false and log error if getting model details fails', async () => { + const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); + esClientMock.ml.getTrainedModels.mockRejectedValue(new Error('error happened')); + const result = await client.isModelInstalled(); + expect(result).toBe(false); + expect(logger.error).toHaveBeenCalled(); + }); + }); + + describe('isInferenceEndpointExists', () => { + it('returns true when the model is fully allocated and started in ESS', async () => { + const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); + esClientMock.ml.getTrainedModelsStats.mockResolvedValueOnce({ + trained_model_stats: [ + { + deployment_stats: { + state: 'started', + // @ts-expect-error not full response interface + allocation_status: { state: 'fully_allocated' }, + }, + }, + ], + }); + + const result = await client.isInferenceEndpointExists(); + + expect(result).toBe(true); + }); + + it('returns true when the model is started in serverless', async () => { + const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); + esClientMock.ml.getTrainedModelsStats.mockResolvedValueOnce({ + trained_model_stats: [ + { + deployment_stats: { + // @ts-expect-error not full response interface + nodes: [{ routing_state: { routing_state: 'started' } }], + }, + }, + ], + }); + + const result = await client.isInferenceEndpointExists(); + + expect(result).toBe(true); + }); + + it('returns false when the model is not fully allocated in ESS', async () => { + const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); + esClientMock.ml.getTrainedModelsStats.mockResolvedValueOnce({ + trained_model_stats: [ + { + deployment_stats: { + state: 'started', + // @ts-expect-error not full response interface + allocation_status: { state: 'partially_allocated' }, + }, + }, + ], + }); + + const result = await client.isInferenceEndpointExists(); + + expect(result).toBe(false); + }); + + it('returns false when the model is not started in serverless', async () => { + const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); + esClientMock.ml.getTrainedModelsStats.mockResolvedValueOnce({ + trained_model_stats: [ + { + deployment_stats: { + // @ts-expect-error not full response interface + nodes: [{ routing_state: { routing_state: 'stopped' } }], + }, + }, + ], + }); + + const result = await client.isInferenceEndpointExists(); + + expect(result).toBe(false); + }); + + it('returns false when an error occurs during the check', async () => { + const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); + esClientMock.ml.getTrainedModelsStats.mockRejectedValueOnce(new Error('Mocked Error')); + + const result = await client.isInferenceEndpointExists(); + + expect(result).toBe(false); + }); + + it('should return false if inference api returns undefined', async () => { + const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); + // @ts-ignore + esClientMock.inference.get.mockResolvedValueOnce(undefined); + const result = await client.isInferenceEndpointExists(); + expect(result).toBe(false); + }); + + it('should return false when inference check throws an error', async () => { + const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); + esClientMock.inference.get.mockRejectedValueOnce(new Error('Mocked Error')); + const result = await client.isInferenceEndpointExists(); + expect(result).toBe(false); + }); + }); + + describe('setupKnowledgeBase', () => { + it('should install, deploy, and load docs if not already done', async () => { + // @ts-expect-error not full response interface + esClientMock.search.mockResolvedValue({}); + + const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); + await client.setupKnowledgeBase({ soClient: savedObjectClient }); + + // install model + expect(trainedModelsProvider).toHaveBeenCalledWith({}, savedObjectClient); + expect(installElasticModel).toHaveBeenCalledWith('elser-id'); + + expect(loadSecurityLabs).toHaveBeenCalled(); + }); + + it('should skip installation and deployment if model is already installed and deployed', async () => { + mockGetSecurityLabsDocsCount.mockResolvedValue(1); + esClientMock.ml.getTrainedModels.mockResolvedValue({ + count: 1, + trained_model_configs: [ + { fully_defined: true, model_id: '', tags: [], input: { field_names: ['content'] } }, + ], + }); + esClientMock.ml.getTrainedModelsStats.mockResolvedValue({ + trained_model_stats: [ + { + deployment_stats: { + state: 'started', + // @ts-expect-error not full response interface + allocation_status: { + state: 'fully_allocated', + }, + }, + }, + ], + }); + const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); + + await client.setupKnowledgeBase({ soClient: savedObjectClient }); + + expect(installElasticModel).not.toHaveBeenCalled(); + expect(esClientMock.ml.startTrainedModelDeployment).not.toHaveBeenCalled(); + expect(loadSecurityLabs).not.toHaveBeenCalled(); + }); + + it('should handle errors during installation and deployment', async () => { + // @ts-expect-error not full response interface + esClientMock.search.mockResolvedValue({}); + esClientMock.ml.getTrainedModels.mockResolvedValue({ + count: 0, + trained_model_configs: [ + { fully_defined: false, model_id: '', tags: [], input: { field_names: ['content'] } }, + ], + }); + mockLoadSecurityLabs.mockRejectedValue(new Error('Installation error')); + const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); + + await expect(client.setupKnowledgeBase({ soClient: savedObjectClient })).rejects.toThrow( + 'Error setting up Knowledge Base: Installation error' + ); + expect(mockOptions.logger.error).toHaveBeenCalledWith( + 'Error setting up Knowledge Base: Installation error' + ); + }); + }); + + describe('addKnowledgeBaseDocuments', () => { + const documents = [ + { + pageContent: 'Document 1', + metadata: { kbResource: 'user', source: 'user', required: false }, + }, + ]; + it('should add documents to the knowledge base', async () => { + esClientMock.bulk.mockResolvedValue({ + items: [ + { + create: { + status: 200, + _id: '123', + _index: 'index', + }, + }, + ], + took: 9999, + errors: false, + }); + const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); + mockOptions.currentUser = mockUser1; + + const result = await client.addKnowledgeBaseDocuments({ documents }); + + expect(result).toHaveLength(1); + expect(result[0]).toEqual(getKnowledgeBaseEntryMock()); + }); + + it('should swallow errors during bulk write', async () => { + esClientMock.bulk.mockRejectedValueOnce(new Error('Bulk write error')); + const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); + mockOptions.currentUser = mockUser1; + + const result = await client.addKnowledgeBaseDocuments({ documents }); + expect(result).toEqual([]); + }); + }); + + describe('isSecurityLabsDocsLoaded', () => { + it('should resolve to true when docs exist', async () => { + const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); + mockOptions.currentUser = mockUser1; + + const results = await client.isSecurityLabsDocsLoaded(); + + expect(results).toEqual(true); + }); + it('should resolve to false when docs do not exist', async () => { + // @ts-expect-error not full response interface + esClientMock.search.mockResolvedValueOnce({ hits: { hits: [] } }); + const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); + mockOptions.currentUser = mockUser1; + + const results = await client.isSecurityLabsDocsLoaded(); + + expect(results).toEqual(false); + }); + it('should resolve to false when docs error', async () => { + esClientMock.search.mockRejectedValueOnce(new Error('Search error')); + const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); + mockOptions.currentUser = mockUser1; + + const results = await client.isSecurityLabsDocsLoaded(); + + expect(results).toEqual(false); + }); + }); + + describe('getKnowledgeBaseDocumentEntries', () => { + it('should fetch documents based on query and filters', async () => { + const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); + mockOptions.currentUser = mockUser1; + + const results = await client.getKnowledgeBaseDocumentEntries({ + query: 'test query', + kbResource: 'security_labs', + }); + + expect(results).toHaveLength(1); + expect(results[0].pageContent).toBe('test'); + expect(results[0].metadata.kbResource).toBe('test'); + }); + + it('should swallow errors during search', async () => { + const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); + mockOptions.currentUser = mockUser1; + + esClientMock.search.mockRejectedValueOnce(new Error('Search error')); + + const results = await client.getKnowledgeBaseDocumentEntries({ + query: 'test query', + }); + expect(results).toEqual([]); + }); + + it('should return an empty array if no documents are found', async () => { + const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); + mockOptions.currentUser = mockUser1; + + // @ts-expect-error not full response interface + esClientMock.search.mockResolvedValueOnce({ hits: { hits: [] } }); + + const results = await client.getKnowledgeBaseDocumentEntries({ + query: 'test query', + }); + + expect(results).toEqual([]); + }); + }); + + describe('getRequiredKnowledgeBaseDocumentEntries', () => { + it('should throw is user is not found', async () => { + const assistantKnowledgeBaseDataClient = new AIAssistantKnowledgeBaseDataClient({ + ...mockOptions, + currentUser: null, + }); + await expect( + assistantKnowledgeBaseDataClient.getRequiredKnowledgeBaseDocumentEntries() + ).rejects.toThrowError( + 'Authenticated user not found! Ensure kbDataClient was initialized from a request.' + ); + }); + it('should fetch the required knowledge base entry successfully', async () => { + const assistantKnowledgeBaseDataClient = new AIAssistantKnowledgeBaseDataClient(mockOptions); + const result = + await assistantKnowledgeBaseDataClient.getRequiredKnowledgeBaseDocumentEntries(); + + expect(esClientMock.search).toHaveBeenCalledTimes(1); + + expect(result).toEqual([ + getKnowledgeBaseEntryMock(getCreateKnowledgeBaseEntrySchemaMock({ required: true })), + ]); + }); + it('should return empty array if unexpected response from findDocuments', async () => { + // @ts-expect-error not full response interface + esClientMock.search.mockResolvedValue({}); + + const assistantKnowledgeBaseDataClient = new AIAssistantKnowledgeBaseDataClient(mockOptions); + const result = + await assistantKnowledgeBaseDataClient.getRequiredKnowledgeBaseDocumentEntries(); + + expect(esClientMock.search).toHaveBeenCalledTimes(1); + + expect(result).toEqual([]); + expect(logger.error).toHaveBeenCalledTimes(2); + }); + }); + + describe('createKnowledgeBaseEntry', () => { + const knowledgeBaseEntry = getCreateKnowledgeBaseEntrySchemaMock(); + it('should create a new Knowledge Base entry', async () => { + const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); + mockOptions.currentUser = mockUser1; + + const result = await client.createKnowledgeBaseEntry({ telemetry, knowledgeBaseEntry }); + expect(result).toEqual(getKnowledgeBaseEntryMock()); + }); + + it('should throw error if user is not authenticated', async () => { + const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); + mockOptions.currentUser = null; + + await expect( + client.createKnowledgeBaseEntry({ telemetry, knowledgeBaseEntry }) + ).rejects.toThrow('Authenticated user not found!'); + }); + + it('should throw error if user lacks privileges to create global entries', async () => { + const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); + mockOptions.currentUser = mockUser1; + mockOptions.manageGlobalKnowledgeBaseAIAssistant = false; + + await expect( + client.createKnowledgeBaseEntry({ telemetry, knowledgeBaseEntry, global: true }) + ).rejects.toThrow('User lacks privileges to create global knowledge base entries'); + }); + }); + + describe('getAssistantTools', () => { + it('should return structured tools for relevant index entries', async () => { + IndexPatternsFetcher.prototype.getExistingIndices = jest.fn().mockResolvedValue(['test']); + esClientMock.search.mockReturnValue( + // @ts-expect-error not full response interface + getKnowledgeBaseEntrySearchEsMock('index') + ); + const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); + mockOptions.currentUser = mockUser1; + + const result = await client.getAssistantTools({ + esClient: esClientMock, + }); + + expect(result).toHaveLength(1); + expect(result[0]).toBeInstanceOf(DynamicStructuredTool); + }); + + it('should return an empty array if no relevant index entries are found', async () => { + const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); + mockOptions.currentUser = mockUser1; + // @ts-expect-error not full response interface + esClientMock.search.mockResolvedValueOnce({ hits: { hits: [] } }); + + const result = await client.getAssistantTools({ + esClient: esClientMock, + }); + + expect(result).toEqual([]); + }); + + it('should swallow errors during fetching index entries', async () => { + const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); + mockOptions.currentUser = mockUser1; + esClientMock.search.mockRejectedValueOnce(new Error('Error fetching index entries')); + + const result = await client.getAssistantTools({ + esClient: esClientMock, + }); + + expect(result).toEqual([]); + }); + }); + + describe('createInferenceEndpoint', () => { + it('should create a new Knowledge Base entry', async () => { + const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); + + esClientMock.inference.put.mockResolvedValueOnce({ + inference_id: 'id', + task_type: 'completion', + service: 'string', + service_settings: {}, + task_settings: {}, + }); + + await client.createInferenceEndpoint(); + + await expect(client.createInferenceEndpoint()).resolves.not.toThrow(); + expect(esClientMock.inference.put).toHaveBeenCalled(); + }); + + it('should throw error if user is not authenticated', async () => { + const client = new AIAssistantKnowledgeBaseDataClient(mockOptions); + + esClientMock.inference.put.mockRejectedValueOnce(new Error('Inference error')); + + await expect(client.createInferenceEndpoint()).rejects.toThrow('Inference error'); + expect(esClientMock.inference.put).toHaveBeenCalled(); + }); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index fae987b6d5083..231aa1c319da4 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -55,7 +55,7 @@ export interface GetAIAssistantKnowledgeBaseDataClientParams { manageGlobalKnowledgeBaseAIAssistant?: boolean; } -interface KnowledgeBaseDataClientParams extends AIAssistantDataClientParams { +export interface KnowledgeBaseDataClientParams extends AIAssistantDataClientParams { ml: MlPluginSetup; getElserId: GetElser; getIsKBSetupInProgress: () => boolean; diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/transforms.test.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/transforms.test.ts new file mode 100644 index 0000000000000..b0451774770b8 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/transforms.test.ts @@ -0,0 +1,64 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { transformESSearchToKnowledgeBaseEntry, transformESToKnowledgeBase } from './transforms'; +import { + getKnowledgeBaseEntrySearchEsMock, + documentEntry, +} from '../../__mocks__/knowledge_base_entry_schema.mock'; + +describe('transforms', () => { + describe('transformESSearchToKnowledgeBaseEntry', () => { + it('should transform Elasticsearch search response to KnowledgeBaseEntryResponse', () => { + const esResponse = getKnowledgeBaseEntrySearchEsMock('document'); + + const result = transformESSearchToKnowledgeBaseEntry(esResponse); + expect(result).toEqual([ + { + id: '1', + createdAt: documentEntry.created_at, + createdBy: documentEntry.created_by, + updatedAt: documentEntry.updated_at, + updatedBy: documentEntry.updated_by, + type: documentEntry.type, + name: documentEntry.name, + namespace: documentEntry.namespace, + kbResource: documentEntry.kb_resource, + source: documentEntry.source, + required: documentEntry.required, + text: documentEntry.text, + users: documentEntry.users, + }, + ]); + }); + }); + + describe('transformESToKnowledgeBase', () => { + it('should transform Elasticsearch response array to KnowledgeBaseEntryResponse array', () => { + const esResponse = [documentEntry]; + + const result = transformESToKnowledgeBase(esResponse); + expect(result).toEqual([ + { + id: documentEntry.id, + createdAt: documentEntry.created_at, + createdBy: documentEntry.created_by, + updatedAt: documentEntry.updated_at, + updatedBy: documentEntry.updated_by, + type: documentEntry.type, + name: documentEntry.name, + namespace: documentEntry.namespace, + kbResource: documentEntry.kb_resource, + source: documentEntry.source, + required: documentEntry.required, + text: documentEntry.text, + users: documentEntry.users, + }, + ]); + }); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.test.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.test.ts index 23a1a55564415..fb3ffe7442c17 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.test.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.test.ts @@ -10,9 +10,9 @@ import { IndicesGetDataStreamResponse } from '@elastic/elasticsearch/lib/api/typ import { errors as EsErrors } from '@elastic/elasticsearch'; import { ReplaySubject, Subject } from 'rxjs'; import { taskManagerMock } from '@kbn/task-manager-plugin/server/mocks'; -import { AuthenticatedUser } from '@kbn/core-security-common'; import { DEFAULT_NAMESPACE_STRING } from '@kbn/core-saved-objects-utils-server'; import { conversationsDataClientMock } from '../__mocks__/data_clients.mock'; +import { authenticatedUser } from '../__mocks__/user'; import { AIAssistantConversationsDataClient } from '../ai_assistant_data_clients/conversations'; import { AIAssistantService, AIAssistantServiceOpts } from '.'; import { retryUntil } from './create_resource_installation_helper.test'; @@ -93,13 +93,7 @@ const getSpaceResourcesInitialized = async ( const conversationsDataClient = conversationsDataClientMock.create(); -const mockUser1 = { - username: 'my_username', - authentication_realm: { - type: 'my_realm_type', - name: 'my_realm_name', - }, -} as AuthenticatedUser; +const mockUser1 = authenticatedUser; describe('AI Assistant Service', () => { let pluginStop$: Subject; diff --git a/x-pack/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.test.ts b/x-pack/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.test.ts index 35653878fa2d2..9e3d9afeb4ba6 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.test.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.test.ts @@ -5,22 +5,17 @@ * 2.0. */ -import type { AuthenticatedUser, ElasticsearchClient, Logger } from '@kbn/core/server'; +import type { ElasticsearchClient, Logger } from '@kbn/core/server'; import { loggingSystemMock, elasticsearchServiceMock } from '@kbn/core/server/mocks'; import { getCreateConversationSchemaMock, getUpdateConversationSchemaMock, } from '../../__mocks__/conversations_schema.mock'; +import { authenticatedUser } from '../../__mocks__/user'; import { DocumentsDataWriter } from './documents_data_writer'; describe('DocumentsDataWriter', () => { - const mockUser1 = { - username: 'my_username', - authentication_realm: { - type: 'my_realm_type', - name: 'my_realm_name', - }, - } as AuthenticatedUser; + const mockUser1 = authenticatedUser; describe('#bulk', () => { let writer: DocumentsDataWriter; let esClientMock: ElasticsearchClient; diff --git a/x-pack/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.ts b/x-pack/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.ts index 08892038a58b7..f065d0a2f8424 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/data_stream/documents_data_writer.ts @@ -22,7 +22,7 @@ export interface BulkOperationError { }; } -interface WriterBulkResponse { +export interface WriterBulkResponse { errors: BulkOperationError[]; docs_created: string[]; docs_deleted: string[]; diff --git a/x-pack/plugins/elastic_assistant/server/routes/anonymization_fields/bulk_actions_route.test.ts b/x-pack/plugins/elastic_assistant/server/routes/anonymization_fields/bulk_actions_route.test.ts index e8055de3b12b9..d3d1302247052 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/anonymization_fields/bulk_actions_route.test.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/anonymization_fields/bulk_actions_route.test.ts @@ -14,7 +14,7 @@ import { getEmptyFindResult, getFindAnonymizationFieldsResultWithSingleHit, } from '../../__mocks__/response'; -import { AuthenticatedUser } from '@kbn/core-security-common'; +import { authenticatedUser } from '../../__mocks__/user'; import { bulkActionAnonymizationFieldsRoute } from './bulk_actions_route'; import { getAnonymizationFieldMock, @@ -28,14 +28,7 @@ describe('Perform bulk action route', () => { let { clients, context } = requestContextMock.createTools(); let logger: ReturnType; const mockAnonymizationField = getAnonymizationFieldMock(getUpdateAnonymizationFieldSchemaMock()); - const mockUser1 = { - profile_uid: 'u_mGBROF_q5bmFCATbLXAcCwKa0k8JvONAwSruelyKA5E_0', - username: 'my_username', - authentication_realm: { - type: 'my_realm_type', - name: 'my_realm_name', - }, - } as AuthenticatedUser; + const mockUser1 = authenticatedUser; beforeEach(async () => { server = serverMock.create(); diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/bulk_actions_route.test.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/bulk_actions_route.test.ts new file mode 100644 index 0000000000000..eb06e34c33219 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/bulk_actions_route.test.ts @@ -0,0 +1,250 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { elasticsearchClientMock } from '@kbn/core-elasticsearch-client-server-mocks'; +import { requestContextMock } from '../../../__mocks__/request_context'; +import { serverMock } from '../../../__mocks__/server'; +import { + getBasicEmptySearchResponse, + getEmptyFindResult, + getFindKnowledgeBaseEntriesResultWithSingleHit, +} from '../../../__mocks__/response'; +import { getBulkActionKnowledgeBaseEntryRequest, requestMock } from '../../../__mocks__/request'; +import { + documentEntry, + getCreateKnowledgeBaseEntrySchemaMock, + getKnowledgeBaseEntryMock, + getQueryKnowledgeBaseEntryParams, + getUpdateKnowledgeBaseEntrySchemaMock, +} from '../../../__mocks__/knowledge_base_entry_schema.mock'; +import { ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_BULK_ACTION } from '@kbn/elastic-assistant-common'; +import { bulkActionKnowledgeBaseEntriesRoute } from './bulk_actions_route'; +import { authenticatedUser } from '../../../__mocks__/user'; + +const date = '2023-03-28T22:27:28.159Z'; +// @ts-ignore +const { kbResource, namespace, ...entrySansResource } = getUpdateKnowledgeBaseEntrySchemaMock('1'); +const { id, ...documentEntrySansId } = documentEntry; + +describe('Bulk actions knowledge base entry route', () => { + let server: ReturnType; + let { clients, context } = requestContextMock.createTools(); + + const mockBulk = jest.fn().mockResolvedValue({ + errors: [], + docs_created: [], + docs_deleted: [], + docs_updated: [], + took: 0, + }); + beforeAll(() => { + jest.useFakeTimers(); + jest.setSystemTime(new Date(date)); + }); + + afterAll(() => { + jest.useRealTimers(); + }); + beforeEach(() => { + jest.clearAllMocks(); + server = serverMock.create(); + ({ clients, context } = requestContextMock.createTools()); + + // @ts-ignore + clients.elasticAssistant.getAIAssistantKnowledgeBaseDataClient.options = { + manageGlobalKnowledgeBaseAIAssistant: true, + }; + + // @ts-ignore + clients.elasticAssistant.getAIAssistantKnowledgeBaseDataClient.getWriter.mockResolvedValue({ + bulk: mockBulk, + }); + + clients.elasticAssistant.getAIAssistantKnowledgeBaseDataClient.findDocuments.mockResolvedValue( + Promise.resolve(getEmptyFindResult()) + ); // no current knowledge base entries + clients.elasticAssistant.getAIAssistantKnowledgeBaseDataClient.createKnowledgeBaseEntry.mockResolvedValue( + getKnowledgeBaseEntryMock(getQueryKnowledgeBaseEntryParams()) + ); // creation succeeds + + context.core.elasticsearch.client.asCurrentUser.search.mockResolvedValue( + elasticsearchClientMock.createSuccessTransportRequestPromise(getBasicEmptySearchResponse()) + ); + bulkActionKnowledgeBaseEntriesRoute(server.router); + }); + + describe('status codes', () => { + test('returns 200 with a knowledge base entry created via AIAssistantKnowledgeBaseDataClient', async () => { + const response = await server.inject( + getBulkActionKnowledgeBaseEntryRequest({ + create: [getCreateKnowledgeBaseEntrySchemaMock()], + }), + requestContextMock.convertContext(context) + ); + + expect(response.status).toEqual(200); + expect(mockBulk).toHaveBeenCalledWith( + expect.objectContaining({ + documentsToCreate: [ + { + ...documentEntrySansId, + '@timestamp': '2023-03-28T22:27:28.159Z', + created_at: '2023-03-28T22:27:28.159Z', + updated_at: '2023-03-28T22:27:28.159Z', + namespace: 'default', + required: false, + }, + ], + authenticatedUser, + }) + ); + }); + test('returns 200 with a knowledge base entry updated via AIAssistantKnowledgeBaseDataClient', async () => { + clients.elasticAssistant.getAIAssistantKnowledgeBaseDataClient.findDocuments.mockResolvedValue( + Promise.resolve(getFindKnowledgeBaseEntriesResultWithSingleHit()) + ); + + const response = await server.inject( + getBulkActionKnowledgeBaseEntryRequest({ + update: [getUpdateKnowledgeBaseEntrySchemaMock('1')], + }), + requestContextMock.convertContext(context) + ); + expect(response.status).toEqual(200); + expect(mockBulk).toHaveBeenCalledWith( + expect.objectContaining({ + documentsToUpdate: [ + { + ...entrySansResource, + required: false, + kb_resource: kbResource, + updated_at: '2023-03-28T22:27:28.159Z', + updated_by: authenticatedUser.profile_uid, + users: [ + { + id: authenticatedUser.profile_uid, + name: authenticatedUser.username, + }, + ], + }, + ], + authenticatedUser, + }) + ); + }); + test('returns 200 with a knowledge base entry deleted via AIAssistantKnowledgeBaseDataClient', async () => { + clients.elasticAssistant.getAIAssistantKnowledgeBaseDataClient.findDocuments.mockResolvedValue( + Promise.resolve(getFindKnowledgeBaseEntriesResultWithSingleHit()) + ); + + const response = await server.inject( + getBulkActionKnowledgeBaseEntryRequest({ + delete: { ids: ['1'] }, + }), + requestContextMock.convertContext(context) + ); + expect(response.status).toEqual(200); + expect(mockBulk).toHaveBeenCalledWith( + expect.objectContaining({ + documentsToDelete: ['1'], + authenticatedUser, + }) + ); + }); + test('handles all three bulk update actions at once', async () => { + clients.elasticAssistant.getAIAssistantKnowledgeBaseDataClient.findDocuments + .mockResolvedValueOnce(Promise.resolve(getEmptyFindResult())) + .mockResolvedValue(Promise.resolve(getFindKnowledgeBaseEntriesResultWithSingleHit())); + const response = await server.inject( + getBulkActionKnowledgeBaseEntryRequest({ + create: [getCreateKnowledgeBaseEntrySchemaMock()], + delete: { ids: ['1'] }, + update: [getUpdateKnowledgeBaseEntrySchemaMock('1')], + }), + requestContextMock.convertContext(context) + ); + expect(response.status).toEqual(200); + expect(mockBulk).toHaveBeenCalledWith( + expect.objectContaining({ + documentsToCreate: [ + { + ...documentEntrySansId, + '@timestamp': '2023-03-28T22:27:28.159Z', + created_at: '2023-03-28T22:27:28.159Z', + updated_at: '2023-03-28T22:27:28.159Z', + namespace: 'default', + required: false, + }, + ], + documentsToUpdate: [ + { + ...entrySansResource, + required: false, + kb_resource: kbResource, + updated_at: '2023-03-28T22:27:28.159Z', + updated_by: authenticatedUser.profile_uid, + users: [ + { + id: authenticatedUser.profile_uid, + name: authenticatedUser.username, + }, + ], + }, + ], + documentsToDelete: ['1'], + authenticatedUser, + }) + ); + }); + test('returns 401 Unauthorized when request context getCurrentUser is not defined', async () => { + context.elasticAssistant.getCurrentUser.mockReturnValueOnce(null); + const response = await server.inject( + getBulkActionKnowledgeBaseEntryRequest({ + create: [getCreateKnowledgeBaseEntrySchemaMock()], + }), + requestContextMock.convertContext(context) + ); + expect(response.status).toEqual(401); + }); + }); + + describe('unhappy paths', () => { + test('catches error if creation throws', async () => { + clients.elasticAssistant.getAIAssistantKnowledgeBaseDataClient.findDocuments.mockImplementation( + async () => { + throw new Error('Test error'); + } + ); + const response = await server.inject( + getBulkActionKnowledgeBaseEntryRequest({ + create: [getCreateKnowledgeBaseEntrySchemaMock()], + }), + requestContextMock.convertContext(context) + ); + expect(response.status).toEqual(500); + expect(response.body).toEqual({ + message: 'Test error', + status_code: 500, + }); + }); + }); + + describe('request validation', () => { + test('disallows wrong name type', async () => { + const request = requestMock.create({ + method: 'post', + path: ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_BULK_ACTION, + body: { + create: [{ ...getCreateKnowledgeBaseEntrySchemaMock(), name: true }], + }, + }); + const result = server.validate(request); + + expect(result.badRequest).toHaveBeenCalled(); + }); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/bulk_actions_route.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/bulk_actions_route.ts index 756e32883ad87..cac334cd73f96 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/bulk_actions_route.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/bulk_actions_route.ts @@ -249,7 +249,6 @@ export const bulkActionKnowledgeBaseEntriesRoute = (router: ElasticAssistantPlug throw new Error(`Could not find documents to ${operation}: ${nonAvailableIds}.`); } }; - await validateDocumentsModification(body.delete?.ids ?? [], 'delete'); await validateDocumentsModification( body.update?.map((entry) => entry.id) ?? [], diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/create_route.test.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/create_route.test.ts new file mode 100644 index 0000000000000..909ca1e5cb6b2 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/create_route.test.ts @@ -0,0 +1,98 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { elasticsearchClientMock } from '@kbn/core-elasticsearch-client-server-mocks'; +import { requestContextMock } from '../../../__mocks__/request_context'; +import { serverMock } from '../../../__mocks__/server'; +import { createKnowledgeBaseEntryRoute } from './create_route'; +import { getBasicEmptySearchResponse, getEmptyFindResult } from '../../../__mocks__/response'; +import { getCreateKnowledgeBaseEntryRequest, requestMock } from '../../../__mocks__/request'; +import { + getCreateKnowledgeBaseEntrySchemaMock, + getKnowledgeBaseEntryMock, + getQueryKnowledgeBaseEntryParams, +} from '../../../__mocks__/knowledge_base_entry_schema.mock'; +import { authenticatedUser } from '../../../__mocks__/user'; +import { ELASTIC_AI_ASSISTANT_CONVERSATIONS_URL } from '@kbn/elastic-assistant-common'; + +describe('Create knowledge base entry route', () => { + let server: ReturnType; + let { clients, context } = requestContextMock.createTools(); + const mockUser1 = authenticatedUser; + + beforeEach(() => { + server = serverMock.create(); + ({ clients, context } = requestContextMock.createTools()); + + clients.elasticAssistant.getAIAssistantKnowledgeBaseDataClient.findDocuments.mockResolvedValue( + Promise.resolve(getEmptyFindResult()) + ); // no current conversations + clients.elasticAssistant.getAIAssistantKnowledgeBaseDataClient.createKnowledgeBaseEntry.mockResolvedValue( + getKnowledgeBaseEntryMock(getQueryKnowledgeBaseEntryParams()) + ); // creation succeeds + + context.core.elasticsearch.client.asCurrentUser.search.mockResolvedValue( + elasticsearchClientMock.createSuccessTransportRequestPromise(getBasicEmptySearchResponse()) + ); + context.elasticAssistant.getCurrentUser.mockReturnValue(mockUser1); + createKnowledgeBaseEntryRoute(server.router); + }); + + describe('status codes', () => { + test('returns 200 with a conversation created via AIAssistantKnowledgeBaseDataClient', async () => { + const response = await server.inject( + getCreateKnowledgeBaseEntryRequest(), + requestContextMock.convertContext(context) + ); + expect(response.status).toEqual(200); + }); + + test('returns 401 Unauthorized when request context getCurrentUser is not defined', async () => { + context.elasticAssistant.getCurrentUser.mockReturnValueOnce(null); + const response = await server.inject( + getCreateKnowledgeBaseEntryRequest(), + requestContextMock.convertContext(context) + ); + expect(response.status).toEqual(401); + }); + }); + + describe('unhappy paths', () => { + test('catches error if creation throws', async () => { + clients.elasticAssistant.getAIAssistantKnowledgeBaseDataClient.createKnowledgeBaseEntry.mockImplementation( + async () => { + throw new Error('Test error'); + } + ); + const response = await server.inject( + getCreateKnowledgeBaseEntryRequest(), + requestContextMock.convertContext(context) + ); + expect(response.status).toEqual(500); + expect(response.body).toEqual({ + message: 'Test error', + status_code: 500, + }); + }); + }); + + describe('request validation', () => { + test('disallows wrong name type', async () => { + const request = requestMock.create({ + method: 'post', + path: ELASTIC_AI_ASSISTANT_CONVERSATIONS_URL, + body: { + ...getCreateKnowledgeBaseEntrySchemaMock(), + name: true, + }, + }); + const result = server.validate(request); + + expect(result.badRequest).toHaveBeenCalled(); + }); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/find_route.test.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/find_route.test.ts new file mode 100644 index 0000000000000..681a3fc2e08fa --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/find_route.test.ts @@ -0,0 +1,111 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { getKnowledgeBaseEntryFindRequest, requestMock } from '../../../__mocks__/request'; +import { ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_FIND } from '@kbn/elastic-assistant-common'; +import { serverMock } from '../../../__mocks__/server'; +import { requestContextMock } from '../../../__mocks__/request_context'; +import { getFindKnowledgeBaseEntriesResultWithSingleHit } from '../../../__mocks__/response'; +import { findKnowledgeBaseEntriesRoute } from './find_route'; +import type { AuthenticatedUser } from '@kbn/core-security-common'; +const mockUser = { + username: 'my_username', + authentication_realm: { + type: 'my_realm_type', + name: 'my_realm_name', + }, +} as AuthenticatedUser; + +describe('Find Knowledge Base Entries route', () => { + let server: ReturnType; + let { clients, context } = requestContextMock.createTools(); + beforeEach(() => { + server = serverMock.create(); + ({ clients, context } = requestContextMock.createTools()); + context.elasticAssistant.getCurrentUser.mockReturnValue(mockUser); + clients.elasticAssistant.getAIAssistantKnowledgeBaseDataClient.findDocuments.mockResolvedValue( + Promise.resolve(getFindKnowledgeBaseEntriesResultWithSingleHit()) + ); + findKnowledgeBaseEntriesRoute(server.router); + }); + + describe('status codes', () => { + test('returns 200', async () => { + const response = await server.inject( + getKnowledgeBaseEntryFindRequest(), + requestContextMock.convertContext(context) + ); + expect(response.status).toEqual(200); + }); + + test('catches error if search throws error', async () => { + clients.elasticAssistant.getAIAssistantKnowledgeBaseDataClient.findDocuments.mockImplementation( + async () => { + throw new Error('Test error'); + } + ); + const response = await server.inject( + getKnowledgeBaseEntryFindRequest(), + requestContextMock.convertContext(context) + ); + expect(response.status).toEqual(500); + expect(response.body).toEqual({ + message: 'Test error', + status_code: 500, + }); + }); + }); + + describe('request validation', () => { + test('allows optional query params', async () => { + const request = requestMock.create({ + method: 'get', + path: ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_FIND, + query: { + page: 2, + per_page: 20, + sort_field: 'title', + fields: ['field1', 'field2'], + }, + }); + const result = server.validate(request); + + expect(result.ok).toHaveBeenCalled(); + }); + + test('disallows invalid sort fields', async () => { + const request = requestMock.create({ + method: 'get', + path: ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_FIND, + query: { + page: 2, + per_page: 20, + sort_field: 'name', + fields: ['field1', 'field2'], + }, + }); + const result = server.validate(request); + + expect(result.badRequest).toHaveBeenCalledWith( + `sort_field: Invalid enum value. Expected 'created_at' | 'is_default' | 'title' | 'updated_at', received 'name'` + ); + }); + + test('ignores unknown query params', async () => { + const request = requestMock.create({ + method: 'get', + path: ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_FIND, + query: { + invalid_value: 'test 1', + }, + }); + const result = server.validate(request); + + expect(result.ok).toHaveBeenCalled(); + }); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/routes/prompts/bulk_actions_route.test.ts b/x-pack/plugins/elastic_assistant/server/routes/prompts/bulk_actions_route.test.ts index 3ae236f12902f..cb3d71b469589 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/prompts/bulk_actions_route.test.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/prompts/bulk_actions_route.test.ts @@ -9,9 +9,9 @@ import { loggingSystemMock } from '@kbn/core/server/mocks'; import { serverMock } from '../../__mocks__/server'; import { requestContextMock } from '../../__mocks__/request_context'; import { getPromptsBulkActionRequest, requestMock } from '../../__mocks__/request'; +import { authenticatedUser } from '../../__mocks__/user'; import { ELASTIC_AI_ASSISTANT_PROMPTS_URL_BULK_ACTION } from '@kbn/elastic-assistant-common'; import { getEmptyFindResult, getFindPromptsResultWithSingleHit } from '../../__mocks__/response'; -import { AuthenticatedUser } from '@kbn/core-security-common'; import { bulkPromptsRoute } from './bulk_actions_route'; import { getCreatePromptSchemaMock, @@ -25,14 +25,7 @@ describe('Perform bulk action route', () => { let { clients, context } = requestContextMock.createTools(); let logger: ReturnType; const mockPrompt = getPromptMock(getUpdatePromptSchemaMock()); - const mockUser1 = { - profile_uid: 'u_mGBROF_q5bmFCATbLXAcCwKa0k8JvONAwSruelyKA5E_0', - username: 'my_username', - authentication_realm: { - type: 'my_realm_type', - name: 'my_realm_name', - }, - } as AuthenticatedUser; + const mockUser1 = authenticatedUser; beforeEach(async () => { server = serverMock.create(); diff --git a/x-pack/plugins/elastic_assistant/server/routes/user_conversations/append_conversation_messages_route.test.ts b/x-pack/plugins/elastic_assistant/server/routes/user_conversations/append_conversation_messages_route.test.ts index 39c25ac6749e9..fb066f1245fe7 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/user_conversations/append_conversation_messages_route.test.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/user_conversations/append_conversation_messages_route.test.ts @@ -14,19 +14,13 @@ import { getQueryConversationParams, getUpdateConversationSchemaMock, } from '../../__mocks__/conversations_schema.mock'; +import { authenticatedUser } from '../../__mocks__/user'; import { appendConversationMessageRoute } from './append_conversation_messages_route'; -import { AuthenticatedUser } from '@kbn/core-security-common'; describe('Append conversation messages route', () => { let server: ReturnType; let { clients, context } = requestContextMock.createTools(); - const mockUser1 = { - username: 'my_username', - authentication_realm: { - type: 'my_realm_type', - name: 'my_realm_name', - }, - } as AuthenticatedUser; + const mockUser1 = authenticatedUser; beforeEach(() => { server = serverMock.create(); diff --git a/x-pack/plugins/elastic_assistant/server/routes/user_conversations/bulk_actions_route.test.ts b/x-pack/plugins/elastic_assistant/server/routes/user_conversations/bulk_actions_route.test.ts index 582651b4cdc9b..d69f53ecaa6c0 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/user_conversations/bulk_actions_route.test.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/user_conversations/bulk_actions_route.test.ts @@ -9,6 +9,7 @@ import { loggingSystemMock } from '@kbn/core/server/mocks'; import { bulkActionConversationsRoute } from './bulk_actions_route'; import { serverMock } from '../../__mocks__/server'; import { requestContextMock } from '../../__mocks__/request_context'; +import { authenticatedUser } from '../../__mocks__/user'; import { getConversationsBulkActionRequest, requestMock } from '../../__mocks__/request'; import { ELASTIC_AI_ASSISTANT_CONVERSATIONS_URL_BULK_ACTION } from '@kbn/elastic-assistant-common'; import { @@ -21,21 +22,13 @@ import { getPerformBulkActionSchemaMock, getUpdateConversationSchemaMock, } from '../../__mocks__/conversations_schema.mock'; -import { AuthenticatedUser } from '@kbn/core-security-common'; describe('Perform bulk action route', () => { let server: ReturnType; let { clients, context } = requestContextMock.createTools(); let logger: ReturnType; const mockConversation = getConversationMock(getUpdateConversationSchemaMock()); - const mockUser1 = { - profile_uid: 'u_mGBROF_q5bmFCATbLXAcCwKa0k8JvONAwSruelyKA5E_0', - username: 'my_username', - authentication_realm: { - type: 'my_realm_type', - name: 'my_realm_name', - }, - } as AuthenticatedUser; + const mockUser1 = authenticatedUser; beforeEach(async () => { server = serverMock.create(); diff --git a/x-pack/plugins/elastic_assistant/server/routes/user_conversations/create_route.test.ts b/x-pack/plugins/elastic_assistant/server/routes/user_conversations/create_route.test.ts index 0659b8d43a38f..378cde4e9bf65 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/user_conversations/create_route.test.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/user_conversations/create_route.test.ts @@ -16,19 +16,13 @@ import { getConversationMock, getQueryConversationParams, } from '../../__mocks__/conversations_schema.mock'; +import { authenticatedUser } from '../../__mocks__/user'; import { ELASTIC_AI_ASSISTANT_CONVERSATIONS_URL } from '@kbn/elastic-assistant-common'; -import { AuthenticatedUser } from '@kbn/core-security-common'; describe('Create conversation route', () => { let server: ReturnType; let { clients, context } = requestContextMock.createTools(); - const mockUser1 = { - username: 'my_username', - authentication_realm: { - type: 'my_realm_type', - name: 'my_realm_name', - }, - } as AuthenticatedUser; + const mockUser1 = authenticatedUser; beforeEach(() => { server = serverMock.create(); diff --git a/x-pack/plugins/elastic_assistant/server/routes/user_conversations/delete_route.test.ts b/x-pack/plugins/elastic_assistant/server/routes/user_conversations/delete_route.test.ts index 128a380c9221a..8edc493c3239f 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/user_conversations/delete_route.test.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/user_conversations/delete_route.test.ts @@ -10,23 +10,16 @@ import { requestContextMock } from '../../__mocks__/request_context'; import { serverMock } from '../../__mocks__/server'; import { deleteConversationRoute } from './delete_route'; import { getDeleteConversationRequest, requestMock } from '../../__mocks__/request'; - +import { authenticatedUser } from '../../__mocks__/user'; import { getConversationMock, getQueryConversationParams, } from '../../__mocks__/conversations_schema.mock'; -import { AuthenticatedUser } from '@kbn/core-security-common'; describe('Delete conversation route', () => { let server: ReturnType; let { clients, context } = requestContextMock.createTools(); - const mockUser1 = { - username: 'my_username', - authentication_realm: { - type: 'my_realm_type', - name: 'my_realm_name', - }, - } as AuthenticatedUser; + const mockUser1 = authenticatedUser; beforeEach(() => { server = serverMock.create(); diff --git a/x-pack/plugins/elastic_assistant/server/routes/user_conversations/read_route.test.ts b/x-pack/plugins/elastic_assistant/server/routes/user_conversations/read_route.test.ts index d2ea1bb5936d3..705054cad9e00 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/user_conversations/read_route.test.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/user_conversations/read_route.test.ts @@ -7,6 +7,7 @@ import { requestContextMock } from '../../__mocks__/request_context'; import { serverMock } from '../../__mocks__/server'; +import { authenticatedUser } from '../../__mocks__/user'; import { readConversationRoute } from './read_route'; import { getConversationReadRequest, requestMock } from '../../__mocks__/request'; import { @@ -14,18 +15,11 @@ import { getQueryConversationParams, } from '../../__mocks__/conversations_schema.mock'; import { ELASTIC_AI_ASSISTANT_CONVERSATIONS_URL_BY_ID } from '@kbn/elastic-assistant-common'; -import { AuthenticatedUser } from '@kbn/core-security-common'; describe('Read conversation route', () => { let server: ReturnType; let { clients, context } = requestContextMock.createTools(); - const mockUser1 = { - username: 'my_username', - authentication_realm: { - type: 'my_realm_type', - name: 'my_realm_name', - }, - } as AuthenticatedUser; + const mockUser1 = authenticatedUser; const myFakeId = '99403909-ca9b-49ba-9d7a-7e5320e68d05'; beforeEach(() => { diff --git a/x-pack/plugins/elastic_assistant/server/routes/user_conversations/update_route.test.ts b/x-pack/plugins/elastic_assistant/server/routes/user_conversations/update_route.test.ts index dd3b7bf1e577d..e8b2a6fcbd507 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/user_conversations/update_route.test.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/user_conversations/update_route.test.ts @@ -13,19 +13,13 @@ import { getQueryConversationParams, getUpdateConversationSchemaMock, } from '../../__mocks__/conversations_schema.mock'; +import { authenticatedUser } from '../../__mocks__/user'; import { updateConversationRoute } from './update_route'; -import { AuthenticatedUser } from '@kbn/core-security-common'; describe('Update conversation route', () => { let server: ReturnType; let { clients, context } = requestContextMock.createTools(); - const mockUser1 = { - username: 'my_username', - authentication_realm: { - type: 'my_realm_type', - name: 'my_realm_name', - }, - } as AuthenticatedUser; + const mockUser1 = authenticatedUser; beforeEach(() => { server = serverMock.create();