diff --git a/x-pack/test/api_integration/apis/ml/annotations/common_jobs.ts b/x-pack/test/api_integration/apis/ml/annotations/common_jobs.ts new file mode 100644 index 0000000000000..873cdc5d71baa --- /dev/null +++ b/x-pack/test/api_integration/apis/ml/annotations/common_jobs.ts @@ -0,0 +1,58 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License; + * you may not use this file except in compliance with the Elastic License. + */ + +import { ANNOTATION_TYPE } from '../../../../../plugins/ml/common/constants/annotations'; +import { Annotation } from '../../../../../plugins/ml/common/types/annotations'; + +export const commonJobConfig = { + description: 'test_job_annotation', + groups: ['farequote', 'automated', 'single-metric'], + analysis_config: { + bucket_span: '15m', + influencers: [], + detectors: [ + { + function: 'mean', + field_name: 'responsetime', + }, + { + function: 'min', + field_name: 'responsetime', + }, + ], + }, + data_description: { time_field: '@timestamp' }, + analysis_limits: { model_memory_limit: '10mb' }, +}; + +export const createJobConfig = (jobId: string) => { + return { ...commonJobConfig, job_id: jobId }; +}; + +export const testSetupJobConfigs = [1, 2, 3, 4].map((num) => ({ + ...commonJobConfig, + job_id: `job_annotation_${num}_${Date.now()}`, + description: `Test annotation ${num}`, +})); +export const jobIds = testSetupJobConfigs.map((j) => j.job_id); + +export const createAnnotationRequestBody = (jobId: string): Partial => { + return { + timestamp: Date.now(), + end_timestamp: Date.now(), + annotation: 'Test annotation', + job_id: jobId, + type: ANNOTATION_TYPE.ANNOTATION, + event: 'user', + detector_index: 1, + partition_field_name: 'airline', + partition_field_value: 'AAL', + }; +}; + +export const testSetupAnnotations = testSetupJobConfigs.map((job) => + createAnnotationRequestBody(job.job_id) +); diff --git a/x-pack/test/api_integration/apis/ml/annotations/create_annotations.ts b/x-pack/test/api_integration/apis/ml/annotations/create_annotations.ts new file mode 100644 index 0000000000000..14ecf1bfe524e --- /dev/null +++ b/x-pack/test/api_integration/apis/ml/annotations/create_annotations.ts @@ -0,0 +1,89 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License; + * you may not use this file except in compliance with the Elastic License. + */ + +import expect from '@kbn/expect'; + +import { FtrProviderContext } from '../../../ftr_provider_context'; +import { COMMON_REQUEST_HEADERS } from '../../../../functional/services/ml/common'; +import { USER } from '../../../../functional/services/ml/security_common'; +import { Annotation } from '../../../../../plugins/ml/common/types/annotations'; +import { createJobConfig, createAnnotationRequestBody } from './common_jobs'; +// eslint-disable-next-line import/no-default-export +export default ({ getService }: FtrProviderContext) => { + const esArchiver = getService('esArchiver'); + const supertest = getService('supertestWithoutAuth'); + const ml = getService('ml'); + + const jobId = `job_annotation_${Date.now()}`; + const testJobConfig = createJobConfig(jobId); + const annotationRequestBody = createAnnotationRequestBody(jobId); + + describe('create_annotations', function () { + before(async () => { + await esArchiver.loadIfNeeded('ml/farequote'); + await ml.testResources.setKibanaTimeZoneToUTC(); + await ml.api.createAnomalyDetectionJob(testJobConfig); + }); + + after(async () => { + await ml.api.cleanMlIndices(); + }); + + it('should successfully create annotations for anomaly job', async () => { + const { body } = await supertest + .put('/api/ml/annotations/index') + .auth(USER.ML_POWERUSER, ml.securityCommon.getPasswordForUser(USER.ML_POWERUSER)) + .set(COMMON_REQUEST_HEADERS) + .send(annotationRequestBody) + .expect(200); + const annotationId = body._id; + + const fetchedAnnotation = await ml.api.getAnnotationById(annotationId); + + expect(fetchedAnnotation).to.not.be(undefined); + + if (fetchedAnnotation) { + Object.keys(annotationRequestBody).forEach((key) => { + const field = key as keyof Annotation; + expect(fetchedAnnotation[field]).to.eql(annotationRequestBody[field]); + }); + } + expect(fetchedAnnotation?.create_username).to.eql(USER.ML_POWERUSER); + }); + + it('should successfully create annotation for user with ML read permissions', async () => { + const { body } = await supertest + .put('/api/ml/annotations/index') + .auth(USER.ML_VIEWER, ml.securityCommon.getPasswordForUser(USER.ML_VIEWER)) + .set(COMMON_REQUEST_HEADERS) + .send(annotationRequestBody) + .expect(200); + + const annotationId = body._id; + const fetchedAnnotation = await ml.api.getAnnotationById(annotationId); + expect(fetchedAnnotation).to.not.be(undefined); + if (fetchedAnnotation) { + Object.keys(annotationRequestBody).forEach((key) => { + const field = key as keyof Annotation; + expect(fetchedAnnotation[field]).to.eql(annotationRequestBody[field]); + }); + } + expect(fetchedAnnotation?.create_username).to.eql(USER.ML_VIEWER); + }); + + it('should not allow to create annotation for unauthorized user', async () => { + const { body } = await supertest + .put('/api/ml/annotations/index') + .auth(USER.ML_UNAUTHORIZED, ml.securityCommon.getPasswordForUser(USER.ML_UNAUTHORIZED)) + .set(COMMON_REQUEST_HEADERS) + .send(annotationRequestBody) + .expect(404); + + expect(body.error).to.eql('Not Found'); + expect(body.message).to.eql('Not Found'); + }); + }); +}; diff --git a/x-pack/test/api_integration/apis/ml/annotations/delete_annotations.ts b/x-pack/test/api_integration/apis/ml/annotations/delete_annotations.ts new file mode 100644 index 0000000000000..4fbb26e9b5a3e --- /dev/null +++ b/x-pack/test/api_integration/apis/ml/annotations/delete_annotations.ts @@ -0,0 +1,91 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License; + * you may not use this file except in compliance with the Elastic License. + */ + +import expect from '@kbn/expect'; +import { FtrProviderContext } from '../../../ftr_provider_context'; +import { COMMON_REQUEST_HEADERS } from '../../../../functional/services/ml/common'; +import { USER } from '../../../../functional/services/ml/security_common'; +import { testSetupJobConfigs, jobIds, testSetupAnnotations } from './common_jobs'; + +// eslint-disable-next-line import/no-default-export +export default ({ getService }: FtrProviderContext) => { + const esArchiver = getService('esArchiver'); + const supertest = getService('supertestWithoutAuth'); + const ml = getService('ml'); + + describe('delete_annotations', function () { + before(async () => { + await esArchiver.loadIfNeeded('ml/farequote'); + await ml.testResources.setKibanaTimeZoneToUTC(); + + // generate one annotation for each job + for (let i = 0; i < testSetupJobConfigs.length; i++) { + const job = testSetupJobConfigs[i]; + const annotationToIndex = testSetupAnnotations[i]; + await ml.api.createAnomalyDetectionJob(job); + await ml.api.indexAnnotation(annotationToIndex); + } + }); + + after(async () => { + await ml.api.cleanMlIndices(); + }); + + it('should delete annotation by id', async () => { + const annotationsForJob = await ml.api.getAnnotations(jobIds[0]); + expect(annotationsForJob).to.have.length(1); + + const annotationIdToDelete = annotationsForJob[0]._id; + + const { body } = await supertest + .delete(`/api/ml/annotations/delete/${annotationIdToDelete}`) + .auth(USER.ML_POWERUSER, ml.securityCommon.getPasswordForUser(USER.ML_POWERUSER)) + .set(COMMON_REQUEST_HEADERS) + .expect(200); + + expect(body._id).to.eql(annotationIdToDelete); + expect(body.result).to.eql('deleted'); + + await ml.api.waitForAnnotationNotToExist(annotationIdToDelete); + }); + + it('should delete annotation by id for user with viewer permission', async () => { + const annotationsForJob = await ml.api.getAnnotations(jobIds[1]); + expect(annotationsForJob).to.have.length(1); + + const annotationIdToDelete = annotationsForJob[0]._id; + + const { body } = await supertest + .delete(`/api/ml/annotations/delete/${annotationIdToDelete}`) + .auth(USER.ML_VIEWER, ml.securityCommon.getPasswordForUser(USER.ML_VIEWER)) + .set(COMMON_REQUEST_HEADERS) + .expect(200); + + expect(body._id).to.eql(annotationIdToDelete); + expect(body.result).to.eql('deleted'); + + await ml.api.waitForAnnotationNotToExist(annotationIdToDelete); + }); + + it('should not delete annotation for unauthorized user', async () => { + const annotationsForJob = await ml.api.getAnnotations(jobIds[2]); + expect(annotationsForJob).to.have.length(1); + + const annotationIdToDelete = annotationsForJob[0]._id; + + const { body } = await supertest + .delete(`/api/ml/annotations/delete/${annotationIdToDelete}`) + .auth(USER.ML_UNAUTHORIZED, ml.securityCommon.getPasswordForUser(USER.ML_UNAUTHORIZED)) + .set(COMMON_REQUEST_HEADERS) + .expect(404); + + expect(body.error).to.eql('Not Found'); + expect(body.message).to.eql('Not Found'); + + await ml.api.waitForAnnotationToExist(annotationIdToDelete); + }); + }); +}; diff --git a/x-pack/test/api_integration/apis/ml/annotations/get_annotations.ts b/x-pack/test/api_integration/apis/ml/annotations/get_annotations.ts new file mode 100644 index 0000000000000..710473eed6901 --- /dev/null +++ b/x-pack/test/api_integration/apis/ml/annotations/get_annotations.ts @@ -0,0 +1,130 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License; + * you may not use this file except in compliance with the Elastic License. + */ + +import expect from '@kbn/expect'; +import { omit } from 'lodash'; +import { FtrProviderContext } from '../../../ftr_provider_context'; +import { COMMON_REQUEST_HEADERS } from '../../../../functional/services/ml/common'; +import { USER } from '../../../../functional/services/ml/security_common'; +import { testSetupJobConfigs, jobIds, testSetupAnnotations } from './common_jobs'; + +// eslint-disable-next-line import/no-default-export +export default ({ getService }: FtrProviderContext) => { + const esArchiver = getService('esArchiver'); + const supertest = getService('supertestWithoutAuth'); + const ml = getService('ml'); + + describe('get_annotations', function () { + before(async () => { + await esArchiver.loadIfNeeded('ml/farequote'); + await ml.testResources.setKibanaTimeZoneToUTC(); + + // generate one annotation for each job + for (let i = 0; i < testSetupJobConfigs.length; i++) { + const job = testSetupJobConfigs[i]; + const annotationToIndex = testSetupAnnotations[i]; + await ml.api.createAnomalyDetectionJob(job); + await ml.api.indexAnnotation(annotationToIndex); + } + }); + + after(async () => { + await ml.api.cleanMlIndices(); + }); + + it('should fetch all annotations for jobId', async () => { + const requestBody = { + jobIds: [jobIds[0]], + earliestMs: 1454804100000, + latestMs: Date.now(), + maxAnnotations: 500, + }; + const { body } = await supertest + .post('/api/ml/annotations') + .auth(USER.ML_POWERUSER, ml.securityCommon.getPasswordForUser(USER.ML_POWERUSER)) + .set(COMMON_REQUEST_HEADERS) + .send(requestBody) + .expect(200); + + expect(body.success).to.eql(true); + expect(body.annotations).not.to.be(undefined); + [jobIds[0]].forEach((jobId, idx) => { + expect(body.annotations).to.have.property(jobId); + expect(body.annotations[jobId]).to.have.length(1); + + const indexedAnnotation = omit(body.annotations[jobId][0], '_id'); + expect(indexedAnnotation).to.eql(testSetupAnnotations[idx]); + }); + }); + + it('should fetch all annotations for multiple jobs', async () => { + const requestBody = { + jobIds, + earliestMs: 1454804100000, + latestMs: Date.now(), + maxAnnotations: 500, + }; + const { body } = await supertest + .post('/api/ml/annotations') + .auth(USER.ML_POWERUSER, ml.securityCommon.getPasswordForUser(USER.ML_POWERUSER)) + .set(COMMON_REQUEST_HEADERS) + .send(requestBody) + .expect(200); + + expect(body.success).to.eql(true); + expect(body.annotations).not.to.be(undefined); + jobIds.forEach((jobId, idx) => { + expect(body.annotations).to.have.property(jobId); + expect(body.annotations[jobId]).to.have.length(1); + + const indexedAnnotation = omit(body.annotations[jobId][0], '_id'); + expect(indexedAnnotation).to.eql(testSetupAnnotations[idx]); + }); + }); + + it('should fetch all annotations for user with ML read permissions', async () => { + const requestBody = { + jobIds: testSetupJobConfigs.map((j) => j.job_id), + earliestMs: 1454804100000, + latestMs: Date.now(), + maxAnnotations: 500, + }; + const { body } = await supertest + .post('/api/ml/annotations') + .auth(USER.ML_VIEWER, ml.securityCommon.getPasswordForUser(USER.ML_VIEWER)) + .set(COMMON_REQUEST_HEADERS) + .send(requestBody) + .expect(200); + expect(body.success).to.eql(true); + expect(body.annotations).not.to.be(undefined); + jobIds.forEach((jobId, idx) => { + expect(body.annotations).to.have.property(jobId); + expect(body.annotations[jobId]).to.have.length(1); + + const indexedAnnotation = omit(body.annotations[jobId][0], '_id'); + expect(indexedAnnotation).to.eql(testSetupAnnotations[idx]); + }); + }); + + it('should not allow to fetch annotation for unauthorized user', async () => { + const requestBody = { + jobIds: testSetupJobConfigs.map((j) => j.job_id), + earliestMs: 1454804100000, + latestMs: Date.now(), + maxAnnotations: 500, + }; + const { body } = await supertest + .post('/api/ml/annotations') + .auth(USER.ML_UNAUTHORIZED, ml.securityCommon.getPasswordForUser(USER.ML_UNAUTHORIZED)) + .set(COMMON_REQUEST_HEADERS) + .send(requestBody) + .expect(404); + + expect(body.error).to.eql('Not Found'); + expect(body.message).to.eql('Not Found'); + }); + }); +}; diff --git a/x-pack/test/api_integration/apis/ml/annotations/index.ts b/x-pack/test/api_integration/apis/ml/annotations/index.ts new file mode 100644 index 0000000000000..7d73ee43d4d99 --- /dev/null +++ b/x-pack/test/api_integration/apis/ml/annotations/index.ts @@ -0,0 +1,16 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License; + * you may not use this file except in compliance with the Elastic License. + */ + +import { FtrProviderContext } from '../../../ftr_provider_context'; + +export default function ({ loadTestFile }: FtrProviderContext) { + describe('annotations', function () { + loadTestFile(require.resolve('./create_annotations')); + loadTestFile(require.resolve('./get_annotations')); + loadTestFile(require.resolve('./delete_annotations')); + loadTestFile(require.resolve('./update_annotations')); + }); +} diff --git a/x-pack/test/api_integration/apis/ml/annotations/update_annotations.ts b/x-pack/test/api_integration/apis/ml/annotations/update_annotations.ts new file mode 100644 index 0000000000000..ba73617151120 --- /dev/null +++ b/x-pack/test/api_integration/apis/ml/annotations/update_annotations.ts @@ -0,0 +1,175 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License; + * you may not use this file except in compliance with the Elastic License. + */ + +import expect from '@kbn/expect'; +import { FtrProviderContext } from '../../../ftr_provider_context'; +import { COMMON_REQUEST_HEADERS } from '../../../../functional/services/ml/common'; +import { USER } from '../../../../functional/services/ml/security_common'; +import { ANNOTATION_TYPE } from '../../../../../plugins/ml/common/constants/annotations'; +import { Annotation } from '../../../../../plugins/ml/common/types/annotations'; +import { testSetupJobConfigs, jobIds, testSetupAnnotations } from './common_jobs'; + +// eslint-disable-next-line import/no-default-export +export default ({ getService }: FtrProviderContext) => { + const esArchiver = getService('esArchiver'); + const supertest = getService('supertestWithoutAuth'); + const ml = getService('ml'); + + const commonAnnotationUpdateRequestBody: Partial = { + timestamp: Date.now(), + end_timestamp: Date.now(), + annotation: 'Updated annotation', + type: ANNOTATION_TYPE.ANNOTATION, + event: 'model_change', + detector_index: 2, + partition_field_name: 'airline', + partition_field_value: 'ANA', + }; + + describe('update_annotations', function () { + before(async () => { + await esArchiver.loadIfNeeded('ml/farequote'); + await ml.testResources.setKibanaTimeZoneToUTC(); + + // generate one annotation for each job + for (let i = 0; i < testSetupJobConfigs.length; i++) { + const job = testSetupJobConfigs[i]; + const annotationToIndex = testSetupAnnotations[i]; + await ml.api.createAnomalyDetectionJob(job); + await ml.api.indexAnnotation(annotationToIndex); + } + }); + + after(async () => { + await ml.api.cleanMlIndices(); + }); + + it('should correctly update annotation by id', async () => { + const annotationsForJob = await ml.api.getAnnotations(jobIds[0]); + expect(annotationsForJob).to.have.length(1); + + const originalAnnotation = annotationsForJob[0]; + const annotationUpdateRequestBody = { + ...commonAnnotationUpdateRequestBody, + job_id: originalAnnotation._source.job_id, + _id: originalAnnotation._id, + }; + + const { body } = await supertest + .put('/api/ml/annotations/index') + .auth(USER.ML_POWERUSER, ml.securityCommon.getPasswordForUser(USER.ML_POWERUSER)) + .set(COMMON_REQUEST_HEADERS) + .send(annotationUpdateRequestBody) + .expect(200); + + expect(body._id).to.eql(originalAnnotation._id); + expect(body.result).to.eql('updated'); + + const updatedAnnotation = await ml.api.getAnnotationById(originalAnnotation._id); + + if (updatedAnnotation) { + Object.keys(commonAnnotationUpdateRequestBody).forEach((key) => { + const field = key as keyof Annotation; + expect(updatedAnnotation[field]).to.eql(annotationUpdateRequestBody[field]); + }); + } + }); + + it('should correctly update annotation for user with viewer permission', async () => { + const annotationsForJob = await ml.api.getAnnotations(jobIds[1]); + expect(annotationsForJob).to.have.length(1); + + const originalAnnotation = annotationsForJob[0]; + const annotationUpdateRequestBody = { + ...commonAnnotationUpdateRequestBody, + job_id: originalAnnotation._source.job_id, + _id: originalAnnotation._id, + }; + + const { body } = await supertest + .put('/api/ml/annotations/index') + .auth(USER.ML_VIEWER, ml.securityCommon.getPasswordForUser(USER.ML_VIEWER)) + .set(COMMON_REQUEST_HEADERS) + .send(annotationUpdateRequestBody) + .expect(200); + + expect(body._id).to.eql(originalAnnotation._id); + expect(body.result).to.eql('updated'); + + const updatedAnnotation = await ml.api.getAnnotationById(originalAnnotation._id); + if (updatedAnnotation) { + Object.keys(commonAnnotationUpdateRequestBody).forEach((key) => { + const field = key as keyof Annotation; + expect(updatedAnnotation[field]).to.eql(annotationUpdateRequestBody[field]); + }); + } + }); + + it('should not update annotation for unauthorized user', async () => { + const annotationsForJob = await ml.api.getAnnotations(jobIds[2]); + expect(annotationsForJob).to.have.length(1); + + const originalAnnotation = annotationsForJob[0]; + + const annotationUpdateRequestBody = { + ...commonAnnotationUpdateRequestBody, + job_id: originalAnnotation._source.job_id, + _id: originalAnnotation._id, + }; + + const { body } = await supertest + .put('/api/ml/annotations/index') + .auth(USER.ML_UNAUTHORIZED, ml.securityCommon.getPasswordForUser(USER.ML_UNAUTHORIZED)) + .set(COMMON_REQUEST_HEADERS) + .send(annotationUpdateRequestBody) + .expect(404); + + expect(body.error).to.eql('Not Found'); + expect(body.message).to.eql('Not Found'); + + const updatedAnnotation = await ml.api.getAnnotationById(originalAnnotation._id); + expect(updatedAnnotation).to.eql(originalAnnotation._source); + }); + + it('should override fields correctly', async () => { + const annotationsForJob = await ml.api.getAnnotations(jobIds[3]); + expect(annotationsForJob).to.have.length(1); + + const originalAnnotation = annotationsForJob[0]; + const annotationUpdateRequestBodyWithMissingFields: Partial = { + timestamp: Date.now(), + end_timestamp: Date.now(), + annotation: 'Updated annotation', + job_id: originalAnnotation._source.job_id, + type: ANNOTATION_TYPE.ANNOTATION, + event: 'model_change', + detector_index: 2, + _id: originalAnnotation._id, + }; + await supertest + .put('/api/ml/annotations/index') + .auth(USER.ML_POWERUSER, ml.securityCommon.getPasswordForUser(USER.ML_POWERUSER)) + .set(COMMON_REQUEST_HEADERS) + .send(annotationUpdateRequestBodyWithMissingFields) + .expect(200); + + const updatedAnnotation = await ml.api.getAnnotationById(originalAnnotation._id); + if (updatedAnnotation) { + Object.keys(annotationUpdateRequestBodyWithMissingFields).forEach((key) => { + if (key !== '_id') { + const field = key as keyof Annotation; + expect(updatedAnnotation[field]).to.eql( + annotationUpdateRequestBodyWithMissingFields[field] + ); + } + }); + } + // validate missing fields in the annotationUpdateRequestBody + expect(updatedAnnotation?.partition_field_name).to.be(undefined); + expect(updatedAnnotation?.partition_field_value).to.be(undefined); + }); + }); +}; diff --git a/x-pack/test/api_integration/apis/ml/index.ts b/x-pack/test/api_integration/apis/ml/index.ts index b29bc47b50394..969f291b0d8b3 100644 --- a/x-pack/test/api_integration/apis/ml/index.ts +++ b/x-pack/test/api_integration/apis/ml/index.ts @@ -60,5 +60,6 @@ export default function ({ getService, loadTestFile }: FtrProviderContext) { loadTestFile(require.resolve('./data_frame_analytics')); loadTestFile(require.resolve('./filters')); loadTestFile(require.resolve('./calendars')); + loadTestFile(require.resolve('./annotations')); }); } diff --git a/x-pack/test/functional/services/ml/api.ts b/x-pack/test/functional/services/ml/api.ts index 9dfec3a17dec0..401a96c5c11bd 100644 --- a/x-pack/test/functional/services/ml/api.ts +++ b/x-pack/test/functional/services/ml/api.ts @@ -5,13 +5,29 @@ */ import expect from '@kbn/expect'; import { ProvidedType } from '@kbn/test/types/ftr'; +import { IndexDocumentParams } from 'elasticsearch'; import { Calendar, CalendarEvent } from '../../../../plugins/ml/server/models/calendar/index'; +import { Annotation } from '../../../../plugins/ml/common/types/annotations'; import { DataFrameAnalyticsConfig } from '../../../../plugins/ml/public/application/data_frame_analytics/common'; import { FtrProviderContext } from '../../ftr_provider_context'; import { DATAFEED_STATE, JOB_STATE } from '../../../../plugins/ml/common/constants/states'; import { DATA_FRAME_TASK_STATE } from '../../../../plugins/ml/public/application/data_frame_analytics/pages/analytics_management/components/analytics_list/common'; import { Datafeed, Job } from '../../../../plugins/ml/common/types/anomaly_detection_jobs'; export type MlApi = ProvidedType; +import { + ML_ANNOTATIONS_INDEX_ALIAS_READ, + ML_ANNOTATIONS_INDEX_ALIAS_WRITE, +} from '../../../../plugins/ml/common/constants/index_patterns'; + +interface EsIndexResult { + _index: string; + _id: string; + _version: number; + result: string; + _shards: any; + _seq_no: number; + _primary_term: number; +} export function MachineLearningAPIProvider({ getService }: FtrProviderContext) { const es = getService('legacyEs'); @@ -634,5 +650,77 @@ export function MachineLearningAPIProvider({ getService }: FtrProviderContext) { } }); }, + + async getAnnotations(jobId: string) { + log.debug(`Fetching annotations for job '${jobId}'...`); + + const results = await es.search({ + index: ML_ANNOTATIONS_INDEX_ALIAS_READ, + body: { + query: { + match: { + job_id: jobId, + }, + }, + }, + }); + expect(results).to.not.be(undefined); + expect(results).to.have.property('hits'); + return results.hits.hits; + }, + + async getAnnotationById(annotationId: string): Promise { + log.debug(`Fetching annotation '${annotationId}'...`); + + const result = await es.search({ + index: ML_ANNOTATIONS_INDEX_ALIAS_READ, + body: { + size: 1, + query: { + match: { + _id: annotationId, + }, + }, + }, + }); + // @ts-ignore due to outdated type for hits.total + if (result.hits.total.value === 1) { + return result?.hits?.hits[0]?._source as Annotation; + } + return undefined; + }, + + async indexAnnotation(annotationRequestBody: Partial) { + log.debug(`Indexing annotation '${JSON.stringify(annotationRequestBody)}'...`); + // @ts-ignore due to outdated type for IndexDocumentParams.type + const params: IndexDocumentParams> = { + index: ML_ANNOTATIONS_INDEX_ALIAS_WRITE, + body: annotationRequestBody, + refresh: 'wait_for', + }; + const results: EsIndexResult = await es.index(params); + await this.waitForAnnotationToExist(results._id); + return results; + }, + + async waitForAnnotationToExist(annotationId: string, errorMsg?: string) { + await retry.tryForTime(30 * 1000, async () => { + if ((await this.getAnnotationById(annotationId)) !== undefined) { + return true; + } else { + throw new Error(errorMsg ?? `annotation '${annotationId}' should exist`); + } + }); + }, + + async waitForAnnotationNotToExist(annotationId: string, errorMsg?: string) { + await retry.tryForTime(30 * 1000, async () => { + if ((await this.getAnnotationById(annotationId)) === undefined) { + return true; + } else { + throw new Error(errorMsg ?? `annotation '${annotationId}' should not exist`); + } + }); + }, }; }