From 9dd86ac9a8d78b0ec66efee052a0667ea7d8228e Mon Sep 17 00:00:00 2001
From: Joshua Li <joshuali925@gmail.com>
Date: Thu, 7 Nov 2024 18:15:44 +0000
Subject: [PATCH] fix flaky unit tests

Signed-off-by: Joshua Li <joshuali925@gmail.com>
---
 .../timefilter/timefilter_service.mock.ts     |  6 +++---
 .../view_components/utils/use_search.test.tsx | 20 ++++++++++++++++---
 2 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/src/plugins/data/public/query/timefilter/timefilter_service.mock.ts b/src/plugins/data/public/query/timefilter/timefilter_service.mock.ts
index 54f98870a658..afc4a18ac89a 100644
--- a/src/plugins/data/public/query/timefilter/timefilter_service.mock.ts
+++ b/src/plugins/data/public/query/timefilter/timefilter_service.mock.ts
@@ -38,10 +38,10 @@ const createSetupContractMock = () => {
     isAutoRefreshSelectorEnabled: jest.fn(),
     isTimeRangeSelectorEnabled: jest.fn(),
     getEnabledUpdated$: jest.fn(),
-    getTimeUpdate$: jest.fn(),
-    getRefreshIntervalUpdate$: jest.fn(),
+    getTimeUpdate$: jest.fn(() => new Observable<unknown>()),
+    getRefreshIntervalUpdate$: jest.fn(() => new Observable<unknown>()),
     getAutoRefreshFetch$: jest.fn(() => new Observable<unknown>()),
-    getFetch$: jest.fn(),
+    getFetch$: jest.fn(() => new Observable<unknown>()),
     getTime: jest.fn(),
     setTime: jest.fn(),
     setRefreshInterval: jest.fn(),
diff --git a/src/plugins/discover/public/application/view_components/utils/use_search.test.tsx b/src/plugins/discover/public/application/view_components/utils/use_search.test.tsx
index 2abb57df6a3d..7f0d95296373 100644
--- a/src/plugins/discover/public/application/view_components/utils/use_search.test.tsx
+++ b/src/plugins/discover/public/application/view_components/utils/use_search.test.tsx
@@ -79,11 +79,17 @@ describe('useSearch', () => {
     const services = createMockServices();
     (services.uiSettings.get as jest.Mock).mockReturnValueOnce(true);
 
-    const { result } = renderHook(() => useSearch(services), { wrapper });
+    const { result, waitForNextUpdate } = renderHook(() => useSearch(services), { wrapper });
 
     expect(result.current.data$.getValue()).toEqual(
       expect.objectContaining({ status: ResultStatus.LOADING })
     );
+
+    // useSearch updates state async in useEffect, wait for it to finish to
+    // avoid warning
+    await act(async () => {
+      await waitForNextUpdate();
+    });
   });
 
   it('should initialize with uninitialized state when search on page load is disabled', async () => {
@@ -94,16 +100,20 @@ describe('useSearch', () => {
       value: 10,
     });
 
-    const { result } = renderHook(() => useSearch(services), { wrapper });
+    const { result, waitForNextUpdate } = renderHook(() => useSearch(services), { wrapper });
     expect(result.current.data$.getValue()).toEqual(
       expect.objectContaining({ status: ResultStatus.UNINITIALIZED })
     );
+
+    await act(async () => {
+      await waitForNextUpdate();
+    });
   });
 
   it('should update startTime when hook rerenders', async () => {
     const services = createMockServices();
 
-    const { result, rerender } = renderHook(() => useSearch(services), {
+    const { result, rerender, waitForNextUpdate } = renderHook(() => useSearch(services), {
       wrapper,
     });
 
@@ -114,6 +124,10 @@ describe('useSearch', () => {
     const newStartTime = result.current.data$.getValue().queryStatus?.startTime;
     expect(newStartTime).toBeDefined();
     expect(newStartTime).not.toEqual(initialStartTime);
+
+    await act(async () => {
+      await waitForNextUpdate();
+    });
   });
 
   it('should reset data observable when dataset changes', async () => {