From 8b1719eae63c4926dbe364465069d0dee44343bb Mon Sep 17 00:00:00 2001
From: Julia Rechkunova <julia.rechkunova@elastic.co>
Date: Tue, 19 Nov 2024 17:22:51 +0100
Subject: [PATCH] [Discover] Update tests

---
 .../extensions/_get_render_app_wrapper.ts     |  4 +++-
 .../extensions/_get_render_app_wrapper.ts     | 19 +++++++++++--------
 2 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/test/functional/apps/discover/context_awareness/extensions/_get_render_app_wrapper.ts b/test/functional/apps/discover/context_awareness/extensions/_get_render_app_wrapper.ts
index 4b08ad76fee5b..b22d54ffe51c5 100644
--- a/test/functional/apps/discover/context_awareness/extensions/_get_render_app_wrapper.ts
+++ b/test/functional/apps/discover/context_awareness/extensions/_get_render_app_wrapper.ts
@@ -13,12 +13,13 @@ import type { WebElementWrapper } from '@kbn/ftr-common-functional-ui-services';
 import type { FtrProviderContext } from '../../ftr_provider_context';
 
 export default function ({ getService, getPageObjects }: FtrProviderContext) {
-  const { common, discover, header, unifiedFieldList, dashboard } = getPageObjects([
+  const { common, discover, header, unifiedFieldList, dashboard, context } = getPageObjects([
     'common',
     'discover',
     'header',
     'unifiedFieldList',
     'dashboard',
+    'context',
   ]);
   const testSubjects = getService('testSubjects');
   const dataViews = getService('dataViews');
@@ -129,6 +130,7 @@ export default function ({ getService, getPageObjects }: FtrProviderContext) {
         await header.waitUntilLoadingHasFinished();
         await browser.refresh();
         await header.waitUntilLoadingHasFinished();
+        await context.waitUntilContextLoadingHasFinished();
 
         await retry.try(async () => {
           messageCell = await dataGrid.getCellElementByColumnName(0, 'message');
diff --git a/x-pack/test_serverless/functional/test_suites/common/discover/context_awareness/extensions/_get_render_app_wrapper.ts b/x-pack/test_serverless/functional/test_suites/common/discover/context_awareness/extensions/_get_render_app_wrapper.ts
index a2dce68cf0805..3506f132f9026 100644
--- a/x-pack/test_serverless/functional/test_suites/common/discover/context_awareness/extensions/_get_render_app_wrapper.ts
+++ b/x-pack/test_serverless/functional/test_suites/common/discover/context_awareness/extensions/_get_render_app_wrapper.ts
@@ -11,14 +11,16 @@ import type { WebElementWrapper } from '@kbn/ftr-common-functional-ui-services';
 import type { FtrProviderContext } from '../../../../../ftr_provider_context';
 
 export default function ({ getService, getPageObjects }: FtrProviderContext) {
-  const { common, discover, header, unifiedFieldList, dashboard, svlCommonPage } = getPageObjects([
-    'common',
-    'discover',
-    'header',
-    'unifiedFieldList',
-    'dashboard',
-    'svlCommonPage',
-  ]);
+  const { common, discover, header, unifiedFieldList, dashboard, context, svlCommonPage } =
+    getPageObjects([
+      'common',
+      'discover',
+      'header',
+      'unifiedFieldList',
+      'dashboard',
+      'context',
+      'svlCommonPage',
+    ]);
   const testSubjects = getService('testSubjects');
   const dataViews = getService('dataViews');
   const dataGrid = getService('dataGrid');
@@ -132,6 +134,7 @@ export default function ({ getService, getPageObjects }: FtrProviderContext) {
         await header.waitUntilLoadingHasFinished();
         await browser.refresh();
         await header.waitUntilLoadingHasFinished();
+        await context.waitUntilContextLoadingHasFinished();
 
         await retry.try(async () => {
           messageCell = await dataGrid.getCellElementByColumnName(0, 'message');