From aec456268091c37b45e0eacd48355ecde072ae00 Mon Sep 17 00:00:00 2001
From: Dima Arnautov <dmitrii.arnautov@elastic.co>
Date: Wed, 30 Nov 2022 16:48:27 +0100
Subject: [PATCH] [ML] Functional tests for the Test Model action  (#146399)

## Summary

Part of #142456

Adds functional tests for the Test model action

### Checklist

- [x] [Unit or functional
tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html)
were updated or added to match the most common scenarios

(cherry picked from commit 9ad78b244abb24eb0dc7f6c15ee511883a4d0ac1)
---
 .../ml_job_editor/ml_job_editor.tsx           |   3 +
 .../inference_input_form/index_input.tsx      |  30 ++--
 .../inference_input_form/text_input.tsx       |  35 +++--
 .../test_models/models/raw_output.tsx         |   2 +-
 .../text_classification/lang_ident_output.tsx |   6 +-
 .../text_classification_output.tsx            |  14 +-
 .../test_models/models/text_input.tsx         |   1 +
 .../test_models/output_loading.tsx            |   4 +-
 .../model_management/model_list.ts            |  17 ++-
 .../test/functional/services/ml/common_ui.ts  |  25 +++-
 x-pack/test/functional/services/ml/index.ts   |   2 +-
 .../functional/services/ml/trained_models.ts  | 138 ++++++++++++++++++
 .../services/ml/trained_models_table.ts       |  36 ++++-
 13 files changed, 266 insertions(+), 47 deletions(-)

diff --git a/x-pack/plugins/ml/public/application/jobs/jobs_list/components/ml_job_editor/ml_job_editor.tsx b/x-pack/plugins/ml/public/application/jobs/jobs_list/components/ml_job_editor/ml_job_editor.tsx
index bbe91f77a7204..f1e4f79f5e7e1 100644
--- a/x-pack/plugins/ml/public/application/jobs/jobs_list/components/ml_job_editor/ml_job_editor.tsx
+++ b/x-pack/plugins/ml/public/application/jobs/jobs_list/components/ml_job_editor/ml_job_editor.tsx
@@ -24,6 +24,7 @@ interface MlJobEditorProps {
   syntaxChecking?: boolean;
   theme?: string;
   onChange?: EuiCodeEditorProps['onChange'];
+  'data-test-subj'?: string;
 }
 export const MLJobEditor: FC<MlJobEditorProps> = ({
   value,
@@ -34,6 +35,7 @@ export const MLJobEditor: FC<MlJobEditorProps> = ({
   syntaxChecking = true,
   theme = 'textmate',
   onChange = () => {},
+  'data-test-subj': dataTestSubj,
 }) => {
   if (mode === ML_EDITOR_MODE.XJSON) {
     try {
@@ -61,6 +63,7 @@ export const MLJobEditor: FC<MlJobEditorProps> = ({
         useSoftTabs: true,
       }}
       onChange={onChange}
+      data-test-subj={dataTestSubj}
     />
   );
 };
diff --git a/x-pack/plugins/ml/public/application/trained_models/models_management/test_models/models/inference_input_form/index_input.tsx b/x-pack/plugins/ml/public/application/trained_models/models_management/test_models/models/inference_input_form/index_input.tsx
index 86ce48851f580..02a9a217f547a 100644
--- a/x-pack/plugins/ml/public/application/trained_models/models_management/test_models/models/inference_input_form/index_input.tsx
+++ b/x-pack/plugins/ml/public/application/trained_models/models_management/test_models/models/inference_input_form/index_input.tsx
@@ -5,7 +5,7 @@
  * 2.0.
  */
 
-import React, { FC, useState, useMemo, useCallback } from 'react';
+import React, { FC, useState, useMemo, useCallback, FormEventHandler } from 'react';
 
 import useObservable from 'react-use/lib/useObservable';
 import { FormattedMessage } from '@kbn/i18n-react';
@@ -18,6 +18,7 @@ import {
   EuiHorizontalRule,
   EuiLoadingSpinner,
   EuiText,
+  EuiForm,
 } from '@elastic/eui';
 
 import { ErrorMessage } from '../../inference_error';
@@ -41,17 +42,21 @@ export const IndexInputForm: FC<Props> = ({ inferrer }) => {
   const outputComponent = useMemo(() => inferrer.getOutputComponent(), [inferrer]);
   const infoComponent = useMemo(() => inferrer.getInfoComponent(), [inferrer]);
 
-  const run = useCallback(async () => {
-    setErrorText(null);
-    try {
-      await inferrer.infer();
-    } catch (e) {
-      setErrorText(extractErrorMessage(e));
-    }
-  }, [inferrer]);
+  const run: FormEventHandler<HTMLFormElement> = useCallback(
+    async (event) => {
+      event.preventDefault();
+      setErrorText(null);
+      try {
+        await inferrer.infer();
+      } catch (e) {
+        setErrorText(extractErrorMessage(e));
+      }
+    },
+    [inferrer]
+  );
 
   return (
-    <>
+    <EuiForm component={'form'} onSubmit={run}>
       <>{infoComponent}</>
       <InferenceInputFormIndexControls inferrer={inferrer} data={data} />
 
@@ -60,9 +65,10 @@ export const IndexInputForm: FC<Props> = ({ inferrer }) => {
       <EuiFlexGroup>
         <EuiFlexItem grow={false}>
           <EuiButton
-            onClick={run}
             disabled={runningState === RUNNING_STATE.RUNNING || isValid === false}
             fullWidth={false}
+            data-test-subj={'mlTestModelTestButton'}
+            type={'submit'}
           >
             <FormattedMessage
               id="xpack.ml.trainedModels.testModelsFlyout.inferenceInputForm.runButton"
@@ -105,6 +111,6 @@ export const IndexInputForm: FC<Props> = ({ inferrer }) => {
         : null}
 
       {runningState === RUNNING_STATE.FINISHED ? <>{outputComponent}</> : null}
-    </>
+    </EuiForm>
   );
 };
diff --git a/x-pack/plugins/ml/public/application/trained_models/models_management/test_models/models/inference_input_form/text_input.tsx b/x-pack/plugins/ml/public/application/trained_models/models_management/test_models/models/inference_input_form/text_input.tsx
index d6ccbf3f3b496..eda62469515ee 100644
--- a/x-pack/plugins/ml/public/application/trained_models/models_management/test_models/models/inference_input_form/text_input.tsx
+++ b/x-pack/plugins/ml/public/application/trained_models/models_management/test_models/models/inference_input_form/text_input.tsx
@@ -5,11 +5,11 @@
  * 2.0.
  */
 
-import React, { FC, useState, useMemo, useCallback } from 'react';
+import React, { FC, useState, useMemo, useCallback, FormEventHandler } from 'react';
 
 import useObservable from 'react-use/lib/useObservable';
 import { FormattedMessage } from '@kbn/i18n-react';
-import { EuiSpacer, EuiButton, EuiTabs, EuiTab } from '@elastic/eui';
+import { EuiSpacer, EuiButton, EuiTabs, EuiTab, EuiForm } from '@elastic/eui';
 
 import { ErrorMessage } from '../../inference_error';
 import { extractErrorMessage } from '../../../../../../../common';
@@ -37,25 +37,30 @@ export const TextInputForm: FC<Props> = ({ inferrer }) => {
   const outputComponent = useMemo(() => inferrer.getOutputComponent(), [inferrer]);
   const infoComponent = useMemo(() => inferrer.getInfoComponent(), [inferrer]);
 
-  const run = useCallback(async () => {
-    setErrorText(null);
-    try {
-      await inferrer.infer();
-    } catch (e) {
-      setErrorText(extractErrorMessage(e));
-    }
-  }, [inferrer]);
+  const run: FormEventHandler<HTMLFormElement> = useCallback(
+    async (event) => {
+      event.preventDefault();
+      setErrorText(null);
+      try {
+        await inferrer.infer();
+      } catch (e) {
+        setErrorText(extractErrorMessage(e));
+      }
+    },
+    [inferrer]
+  );
 
   return (
-    <>
+    <EuiForm component={'form'} onSubmit={run}>
       <>{infoComponent}</>
       <>{inputComponent}</>
       <EuiSpacer size="m" />
       <div>
         <EuiButton
-          onClick={run}
           disabled={runningState === RUNNING_STATE.RUNNING || isValid === false}
           fullWidth={false}
+          data-test-subj={'mlTestModelTestButton'}
+          type="submit"
         >
           <FormattedMessage
             id="xpack.ml.trainedModels.testModelsFlyout.inferenceInputForm.runButton"
@@ -100,13 +105,15 @@ export const TextInputForm: FC<Props> = ({ inferrer }) => {
                 </>
               ) : null}
 
-              {runningState === RUNNING_STATE.FINISHED ? <>{outputComponent}</> : null}
+              {runningState === RUNNING_STATE.FINISHED ? (
+                <div data-test-subj={'mlTestModelOutput'}>{outputComponent}</div>
+              ) : null}
             </>
           ) : (
             <RawOutput inferrer={inferrer} />
           )}
         </>
       ) : null}
-    </>
+    </EuiForm>
   );
 };
diff --git a/x-pack/plugins/ml/public/application/trained_models/models_management/test_models/models/raw_output.tsx b/x-pack/plugins/ml/public/application/trained_models/models_management/test_models/models/raw_output.tsx
index 9b88ed79bfd3c..6656333f259f8 100644
--- a/x-pack/plugins/ml/public/application/trained_models/models_management/test_models/models/raw_output.tsx
+++ b/x-pack/plugins/ml/public/application/trained_models/models_management/test_models/models/raw_output.tsx
@@ -59,7 +59,7 @@ export const RawOutput: FC<{
 
   return (
     <>
-      <MLJobEditor value={rawResponse ?? ''} readOnly={true} />
+      <MLJobEditor data-test-subj={'mlTestModelRawOutput'} value={rawResponse ?? ''} readOnly />
     </>
   );
 };
diff --git a/x-pack/plugins/ml/public/application/trained_models/models_management/test_models/models/text_classification/lang_ident_output.tsx b/x-pack/plugins/ml/public/application/trained_models/models_management/test_models/models/text_classification/lang_ident_output.tsx
index 1d1582629832c..d07e8f8717cd3 100644
--- a/x-pack/plugins/ml/public/application/trained_models/models_management/test_models/models/text_classification/lang_ident_output.tsx
+++ b/x-pack/plugins/ml/public/application/trained_models/models_management/test_models/models/text_classification/lang_ident_output.tsx
@@ -57,10 +57,12 @@ const LanguageIdent: FC<{
 
   return (
     <>
-      <EuiText size="s">{inputText}</EuiText>
+      <EuiText size="s" data-test-subj={'mlTestModelLangIdentInputText'}>
+        {inputText}
+      </EuiText>
       <EuiSpacer size="s" />
       <EuiTitle size="xxs">
-        <h4>{title}</h4>
+        <h4 data-test-subj={'mlTestModelLangIdentTitle'}>{title}</h4>
       </EuiTitle>
 
       <EuiSpacer />
diff --git a/x-pack/plugins/ml/public/application/trained_models/models_management/test_models/models/text_classification/text_classification_output.tsx b/x-pack/plugins/ml/public/application/trained_models/models_management/test_models/models/text_classification/text_classification_output.tsx
index acf556be6009c..ca8a398d48d78 100644
--- a/x-pack/plugins/ml/public/application/trained_models/models_management/test_models/models/text_classification/text_classification_output.tsx
+++ b/x-pack/plugins/ml/public/application/trained_models/models_management/test_models/models/text_classification/text_classification_output.tsx
@@ -5,7 +5,7 @@
  * 2.0.
  */
 
-import React, { FC } from 'react';
+import React, { type FC, Fragment } from 'react';
 import useObservable from 'react-use/lib/useObservable';
 import {
   EuiFlexGroup,
@@ -72,17 +72,17 @@ export const PredictionProbabilityList: FC<{
       ) : null}
 
       {response.map(({ value, predictionProbability }) => (
-        <>
+        <Fragment key={value}>
           <EuiProgress value={predictionProbability * 100} max={100} size="m" />
           <EuiSpacer size="s" />
           <EuiFlexGroup>
-            <>
-              <EuiFlexItem>{value}</EuiFlexItem>
-              <EuiFlexItem grow={false}>{predictionProbability}</EuiFlexItem>
-            </>
+            <EuiFlexItem data-test-subj={`mlTestModelLangIdentInputValue`}>{value}</EuiFlexItem>
+            <EuiFlexItem data-test-subj={`mlTestModelLangIdentInputProbability`} grow={false}>
+              {predictionProbability}
+            </EuiFlexItem>
           </EuiFlexGroup>
           <EuiSpacer />
-        </>
+        </Fragment>
       ))}
     </>
   );
diff --git a/x-pack/plugins/ml/public/application/trained_models/models_management/test_models/models/text_input.tsx b/x-pack/plugins/ml/public/application/trained_models/models_management/test_models/models/text_input.tsx
index 830635af4368e..ead2a0835a210 100644
--- a/x-pack/plugins/ml/public/application/trained_models/models_management/test_models/models/text_input.tsx
+++ b/x-pack/plugins/ml/public/application/trained_models/models_management/test_models/models/text_input.tsx
@@ -44,6 +44,7 @@ export const TextInput: FC<{
         onChange={(e) => {
           setInputText(e.target.value);
         }}
+        data-test-subj={`mlTestModelInputText`}
       />
     </EuiFormRow>
   );
diff --git a/x-pack/plugins/ml/public/application/trained_models/models_management/test_models/output_loading.tsx b/x-pack/plugins/ml/public/application/trained_models/models_management/test_models/output_loading.tsx
index 4cceed23edd25..9c21642acd4b5 100644
--- a/x-pack/plugins/ml/public/application/trained_models/models_management/test_models/output_loading.tsx
+++ b/x-pack/plugins/ml/public/application/trained_models/models_management/test_models/output_loading.tsx
@@ -13,5 +13,7 @@ export const OutputLoadingContent: FC<{ text: string }> = ({ text }) => {
   const actualLines = text.split(/\r\n|\r|\n/).length + 1;
   const lines = actualLines > 4 && actualLines <= 10 ? actualLines : 4;
 
-  return <EuiLoadingContent lines={lines as LineRange} />;
+  return (
+    <EuiLoadingContent data-test-subj={'mlTestModelLoadingContent'} lines={lines as LineRange} />
+  );
 };
diff --git a/x-pack/test/functional/apps/ml/short_tests/model_management/model_list.ts b/x-pack/test/functional/apps/ml/short_tests/model_management/model_list.ts
index 2743df8a81949..04a23cd33e1c7 100644
--- a/x-pack/test/functional/apps/ml/short_tests/model_management/model_list.ts
+++ b/x-pack/test/functional/apps/ml/short_tests/model_management/model_list.ts
@@ -93,7 +93,7 @@ export default function ({ getService }: FtrProviderContext) {
         await ml.trainedModelsTable.assertPipelinesTabContent(false);
       });
 
-      it('displays the built-in model and no actions are enabled', async () => {
+      it('displays the built-in model with only Test action enabled', async () => {
         await ml.testExecution.logTestStep('should display the model in the table');
         await ml.trainedModelsTable.filterWithSearchString(builtInModelData.modelId, 1);
 
@@ -121,6 +121,21 @@ export default function ({ getService }: FtrProviderContext) {
           builtInModelData.modelId,
           false
         );
+
+        await ml.testExecution.logTestStep('should have enabled the button that opens Test flyout');
+        await ml.trainedModelsTable.assertModelTestButtonExists(builtInModelData.modelId, true);
+
+        await ml.trainedModelsTable.testModel(
+          'lang_ident',
+          builtInModelData.modelId,
+          {
+            inputText: 'Goedemorgen! Ik ben een appel.',
+          },
+          {
+            title: 'This looks like Dutch,Flemish',
+            topLang: { code: 'nl', minProbability: 0.9 },
+          }
+        );
       });
 
       it('displays a model with an ingest pipeline and delete action is disabled', async () => {
diff --git a/x-pack/test/functional/services/ml/common_ui.ts b/x-pack/test/functional/services/ml/common_ui.ts
index ef69f909437c1..138450e56a426 100644
--- a/x-pack/test/functional/services/ml/common_ui.ts
+++ b/x-pack/test/functional/services/ml/common_ui.ts
@@ -402,17 +402,28 @@ export function MachineLearningCommonUIProvider({
       });
     },
 
-    async invokeTableRowAction(rowSelector: string, actionTestSubject: string) {
+    async invokeTableRowAction(
+      rowSelector: string,
+      actionTestSubject: string,
+      fromContextMenu: boolean = true
+    ) {
       await retry.tryForTime(30 * 1000, async () => {
-        await this.ensureAllMenuPopoversClosed();
-        await testSubjects.click(`${rowSelector} > euiCollapsedItemActionsButton`);
-        await find.existsByCssSelector('euiContextMenuPanel');
+        if (fromContextMenu) {
+          await this.ensureAllMenuPopoversClosed();
+
+          await testSubjects.click(`${rowSelector} > euiCollapsedItemActionsButton`);
+          await find.existsByCssSelector('euiContextMenuPanel');
 
-        const isEnabled = await testSubjects.isEnabled(actionTestSubject);
+          const isEnabled = await testSubjects.isEnabled(actionTestSubject);
 
-        expect(isEnabled).to.eql(true, `Expected action "${actionTestSubject}" to be enabled.`);
+          expect(isEnabled).to.eql(true, `Expected action "${actionTestSubject}" to be enabled.`);
 
-        await testSubjects.click(actionTestSubject);
+          await testSubjects.click(actionTestSubject);
+        } else {
+          const isEnabled = await testSubjects.isEnabled(`${rowSelector} > ${actionTestSubject}`);
+          expect(isEnabled).to.eql(true, `Expected action "${actionTestSubject}" to be enabled.`);
+          await testSubjects.click(`${rowSelector} > ${actionTestSubject}`);
+        }
       });
     },
   };
diff --git a/x-pack/test/functional/services/ml/index.ts b/x-pack/test/functional/services/ml/index.ts
index 9452baa324898..e360bfecf3a32 100644
--- a/x-pack/test/functional/services/ml/index.ts
+++ b/x-pack/test/functional/services/ml/index.ts
@@ -131,7 +131,7 @@ export function MachineLearningProvider(context: FtrProviderContext) {
   const alerting = MachineLearningAlertingProvider(context, api, commonUI);
   const swimLane = SwimLaneProvider(context);
   const trainedModels = TrainedModelsProvider(context, commonUI);
-  const trainedModelsTable = TrainedModelsTableProvider(context, commonUI);
+  const trainedModelsTable = TrainedModelsTableProvider(context, commonUI, trainedModels);
   const mlNodesPanel = MlNodesPanelProvider(context);
   const notifications = NotificationsProvider(context, commonUI, tableService);
 
diff --git a/x-pack/test/functional/services/ml/trained_models.ts b/x-pack/test/functional/services/ml/trained_models.ts
index c0fb549f97864..09503f7415ed4 100644
--- a/x-pack/test/functional/services/ml/trained_models.ts
+++ b/x-pack/test/functional/services/ml/trained_models.ts
@@ -5,13 +5,76 @@
  * 2.0.
  */
 
+// eslint-disable-next-line max-classes-per-file
 import expect from '@kbn/expect';
+import { ProvidedType } from '@kbn/test';
 import { FtrProviderContext } from '../../ftr_provider_context';
 import { MlCommonUI } from './common_ui';
 
+export type TrainedModelsActions = ProvidedType<typeof TrainedModelsProvider>;
+
+export type ModelType = 'lang_ident';
+
+export interface MappedInputParams {
+  lang_ident: LangIdentInput;
+}
+
+export interface MappedOutput {
+  lang_ident: LangIdentOutput;
+}
+
 export function TrainedModelsProvider({ getService }: FtrProviderContext, mlCommonUI: MlCommonUI) {
   const testSubjects = getService('testSubjects');
   const retry = getService('retry');
+  const browser = getService('browser');
+
+  class TestModelFactory {
+    public static createAssertionInstance(modelType: ModelType) {
+      switch (modelType) {
+        case 'lang_ident':
+          return new TestLangIdentModel();
+        default:
+          throw new Error(`Testing class for ${modelType} is not implemented`);
+      }
+    }
+  }
+
+  class TestModelBase implements TestTrainedModel<BaseInput, unknown> {
+    async setRequiredInput(input: BaseInput): Promise<void> {
+      await testSubjects.setValue('mlTestModelInputText', input.inputText);
+      await this.assertTestInputText(input.inputText);
+    }
+
+    async assertTestInputText(expectedText: string) {
+      const actualValue = await testSubjects.getAttribute('mlTestModelInputText', 'value');
+      expect(actualValue).to.eql(
+        expectedText,
+        `Expected input text to equal ${expectedText}, got ${actualValue}`
+      );
+    }
+
+    assertModelOutput(expectedOutput: unknown): Promise<void> {
+      throw new Error('assertModelOutput has to be implemented per model type');
+    }
+  }
+
+  class TestLangIdentModel
+    extends TestModelBase
+    implements TestTrainedModel<LangIdentInput, LangIdentOutput>
+  {
+    async assertModelOutput(expectedOutput: LangIdentOutput) {
+      const title = await testSubjects.getVisibleText('mlTestModelLangIdentTitle');
+      expect(title).to.eql(expectedOutput.title);
+
+      const values = await testSubjects.findAll('mlTestModelLangIdentInputValue');
+      const topValue = await values[0].getVisibleText();
+      expect(topValue).to.eql(expectedOutput.topLang.code);
+
+      const probabilities = await testSubjects.findAll('mlTestModelLangIdentInputProbability');
+      const topProbability = Number(await probabilities[0].getVisibleText());
+      expect(topProbability).to.above(expectedOutput.topLang.minProbability);
+    }
+  }
 
   return {
     async assertStats(expectedTotalCount: number) {
@@ -28,5 +91,80 @@ export function TrainedModelsProvider({ getService }: FtrProviderContext, mlComm
     async assertRowsNumberPerPage(rowsNumber: 10 | 25 | 100) {
       await mlCommonUI.assertRowsNumberPerPage('mlModelsTableContainer', rowsNumber);
     },
+
+    async assertTestButtonEnabled(expectedValue: boolean = false) {
+      const isEnabled = await testSubjects.isEnabled('mlTestModelTestButton');
+      expect(isEnabled).to.eql(
+        expectedValue,
+        `Expected trained model "Test" button to be '${
+          expectedValue ? 'enabled' : 'disabled'
+        }' (got '${isEnabled ? 'enabled' : 'disabled'}')`
+      );
+    },
+
+    async testModel() {
+      await testSubjects.click('mlTestModelTestButton');
+    },
+
+    async assertTestInputText(expectedText: string) {
+      const actualValue = await testSubjects.getAttribute('mlTestModelInputText', 'value');
+      expect(actualValue).to.eql(
+        expectedText,
+        `Expected input text to equal ${expectedText}, got ${actualValue}`
+      );
+    },
+
+    async waitForResultsToLoad() {
+      await testSubjects.waitForEnabled('mlTestModelTestButton');
+      await retry.tryForTime(5000, async () => {
+        await testSubjects.existOrFail(`mlTestModelOutput`);
+      });
+    },
+
+    async testModelOutput(
+      modelType: ModelType,
+      inputParams: MappedInputParams[typeof modelType],
+      expectedOutput: MappedOutput[typeof modelType]
+    ) {
+      await this.assertTestButtonEnabled(false);
+
+      const modelTest = TestModelFactory.createAssertionInstance(modelType);
+      await modelTest.setRequiredInput(inputParams);
+
+      await this.assertTestButtonEnabled(true);
+      await this.testModel();
+      await this.waitForResultsToLoad();
+
+      await modelTest.assertModelOutput(expectedOutput);
+
+      await this.ensureTestFlyoutClosed();
+    },
+
+    async ensureTestFlyoutClosed() {
+      await retry.tryForTime(5000, async () => {
+        await browser.pressKeys(browser.keys.ESCAPE);
+        await testSubjects.missingOrFail('mlTestModelsFlyout');
+      });
+    },
   };
 }
+
+export interface BaseInput {
+  inputText: string;
+}
+
+export type LangIdentInput = BaseInput;
+
+export interface LangIdentOutput {
+  title: string;
+  topLang: { code: string; minProbability: number };
+}
+
+/**
+ * Interface that needed to be implemented by all model types
+ */
+interface TestTrainedModel<Input extends BaseInput, Output> {
+  setRequiredInput(input: Input): Promise<void>;
+  assertTestInputText(inputText: Input['inputText']): Promise<void>;
+  assertModelOutput(expectedOutput: Output): Promise<void>;
+}
diff --git a/x-pack/test/functional/services/ml/trained_models_table.ts b/x-pack/test/functional/services/ml/trained_models_table.ts
index 6f2b76be1bdb1..8179dc7a5986b 100644
--- a/x-pack/test/functional/services/ml/trained_models_table.ts
+++ b/x-pack/test/functional/services/ml/trained_models_table.ts
@@ -12,6 +12,7 @@ import { upperFirst } from 'lodash';
 import { WebElementWrapper } from '../../../../../test/functional/services/lib/web_element_wrapper';
 import type { FtrProviderContext } from '../../ftr_provider_context';
 import type { MlCommonUI } from './common_ui';
+import { MappedInputParams, MappedOutput, ModelType, TrainedModelsActions } from './trained_models';
 
 export interface TrainedModelRowData {
   id: string;
@@ -23,7 +24,8 @@ export type MlTrainedModelsTable = ProvidedType<typeof TrainedModelsTableProvide
 
 export function TrainedModelsTableProvider(
   { getService }: FtrProviderContext,
-  mlCommonUI: MlCommonUI
+  mlCommonUI: MlCommonUI,
+  trainedModelsActions: TrainedModelsActions
 ) {
   const testSubjects = getService('testSubjects');
   const retry = getService('retry');
@@ -180,6 +182,34 @@ export function TrainedModelsTableProvider(
       );
     }
 
+    public async assertModelTestButtonExists(modelId: string, expectedValue: boolean) {
+      const actionExists = await testSubjects.exists(
+        this.rowSelector(modelId, 'mlModelsTableRowTestAction')
+      );
+      expect(actionExists).to.eql(
+        expectedValue,
+        `Expected test action button for trained model '${modelId}' to be ${
+          expectedValue ? 'visible' : 'hidden'
+        } (got ${actionExists ? 'visible' : 'hidden'})`
+      );
+    }
+
+    public async testModel(
+      modelType: ModelType,
+      modelId: string,
+      inputParams: MappedInputParams[typeof modelType],
+      expectedResult: MappedOutput[typeof modelType]
+    ) {
+      await mlCommonUI.invokeTableRowAction(
+        this.rowSelector(modelId),
+        'mlModelsTableRowTestAction',
+        false
+      );
+      await this.assertTestFlyoutExists();
+
+      await trainedModelsActions.testModelOutput(modelType, inputParams, expectedResult);
+    }
+
     public async deleteModel(modelId: string) {
       await mlCommonUI.invokeTableRowAction(
         this.rowSelector(modelId),
@@ -208,6 +238,10 @@ export function TrainedModelsTableProvider(
       await testSubjects.existOrFail('mlModelsDeleteModal', { timeout: 60 * 1000 });
     }
 
+    public async assertTestFlyoutExists() {
+      await testSubjects.existOrFail('mlTestModelsFlyout', { timeout: 60 * 1000 });
+    }
+
     public async assertStartDeploymentModalExists(expectExist = true) {
       if (expectExist) {
         await testSubjects.existOrFail('mlModelsStartDeploymentModal', { timeout: 60 * 1000 });