[Security Solution] Fix flaky unified component test #189791

x-pack/plugins/security_solution/public/timelines/components/timeline/tabs/query/query_tab_unified_components.test.tsx

@@ -297,21 +297,27 @@ describe('query tab with unified timeline', () => {
     );
   });
 
-  // FLAKY: https://github.com/elastic/kibana/issues/189791
-  describe.skip('pagination', () => {
+  describe('pagination', () => {
     beforeEach(() => {
-      // should return all the records instead just 3
-      // as the case in the default mock
+      // pagination tests need more than 1 record so here
+      // we return 5 records instead of just 1.
       useTimelineEventsMock = jest.fn(() => [
         false,
         {
-          events: structuredClone(mockTimelineData),
+          events: structuredClone(mockTimelineData.slice(0, 5)),
           pageInfo: {
             activePage: 0,
-            totalPages: 10,
+            totalPages: 5,
           },
           refreshedAt: Date.now(),
-          totalCount: 70,
+          /*
+           * `totalCount` could be any number w.r.t this test
+           * and actually means total hits on elastic search
+           * and not the fecthed number of records.
+           *
+           * This helps in testing `sampleSize` and `loadMore`
+           */
+          totalCount: 50,
           loadPage: loadPageMock,
         },
       ]);
@@ -326,21 +332,48 @@ describe('query tab with unified timeline', () => {
     it(
       'should paginate correctly',
       async () => {
-        renderTestComponents();
+        const mockStateWithNoteInTimeline = {
+          ...mockGlobalState,
+          timeline: {
+            ...mockGlobalState.timeline,
+            timelineById: {
+              [TimelineId.test]: {
+                ...mockGlobalState.timeline.timelineById[TimelineId.test],
+                /* 1 record for each page */
+                itemsPerPage: 1,
+                itemsPerPageOptions: [1, 2, 3, 4, 5],
+                savedObjectId: 'timeline-1', // match timelineId in mocked notes data
+                pinnedEventIds: { '1': true },
+              },
+            },
+          },
+        };
 
-        await waitFor(() => {
-          expect(screen.getByTestId('tablePaginationPopoverButton')).toHaveTextContent(
-            'Rows per page: 5'
-          );
-        });
+        render(
+          <TestProviders
+            store={createMockStore({
+              ...structuredClone(mockStateWithNoteInTimeline),
+            })}
+          >
+            <TestComponent />
+          </TestProviders>
+        );
+
+        expect(await screen.findByTestId('discoverDocTable')).toBeVisible();
+        expect(screen.getByTestId('pagination-button-previous')).toBeVisible();
+
+        expect(screen.getByTestId('tablePaginationPopoverButton')).toHaveTextContent(
+          'Rows per page: 1'
+        );
 
         expect(screen.getByTestId('pagination-button-0')).toHaveAttribute('aria-current', 'true');
-        expect(screen.getByTestId('pagination-button-6')).toBeVisible();
+        expect(screen.getByTestId('pagination-button-4')).toBeVisible();
+        expect(screen.queryByTestId('pagination-button-5')).toBeNull();
 
-        fireEvent.click(screen.getByTestId('pagination-button-6'));
+        fireEvent.click(screen.getByTestId('pagination-button-4'));
 
         await waitFor(() => {
-          expect(screen.getByTestId('pagination-button-6')).toHaveAttribute('aria-current', 'true');
+          expect(screen.getByTestId('pagination-button-4')).toHaveAttribute('aria-current', 'true');
         });
       },
       SPECIAL_TEST_TIMEOUT
@@ -349,13 +382,45 @@ describe('query tab with unified timeline', () => {
     it(
       'should load more records according to sample size correctly',
       async () => {
-        renderTestComponents();
+        const mockStateWithNoteInTimeline = {
+          ...mockGlobalState,
+          timeline: {
+            ...mockGlobalState.timeline,
+            timelineById: {
+              [TimelineId.test]: {
+                ...mockGlobalState.timeline.timelineById[TimelineId.test],
+                itemsPerPage: 1,
+                /*
+                 * `sampleSize` is the max number of records that are fetched from elasticsearch
+                 * in one request. If hits > sampleSize, you can fetch more records ( <= sampleSize)
+                 */
+                sampleSize: 5,
+                itemsPerPageOptions: [1, 2, 3, 4, 5],
+                savedObjectId: 'timeline-1', // match timelineId in mocked notes data
+                pinnedEventIds: { '1': true },
+              },
+            },
+          },
+        };
+
+        render(
+          <TestProviders
+            store={createMockStore({
+              ...structuredClone(mockStateWithNoteInTimeline),
+            })}
+          >
+            <TestComponent />
+          </TestProviders>
+        );
+
+        expect(await screen.findByTestId('discoverDocTable')).toBeVisible();
+
         await waitFor(() => {
           expect(screen.getByTestId('pagination-button-0')).toHaveAttribute('aria-current', 'true');
-          expect(screen.getByTestId('pagination-button-6')).toBeVisible();
+          expect(screen.getByTestId('pagination-button-4')).toBeVisible();
         });
         // Go to last page
-        fireEvent.click(screen.getByTestId('pagination-button-6'));
+        fireEvent.click(screen.getByTestId('pagination-button-4'));
         await waitFor(() => {
           expect(screen.getByTestId('dscGridSampleSizeFetchMoreLink')).toBeVisible();
         });