[Security Solution] Unified Timeline - Fix Flaky tests

packages/kbn-unified-data-table/src/utils/get_render_cell_popover.tsx

@@ -19,7 +19,9 @@ import React, { memo, useEffect } from 'react';
  *
  * */
 export const getCustomCellPopoverRenderer = () => {
-  return memo(function RenderCustomCellPopover(props: EuiDataGridCellPopoverElementProps) {
+  const RenderCustomCellPopoverMemoized = memo(function RenderCustomCellPopoverMemoized(
+    props: EuiDataGridCellPopoverElementProps
+  ) {
     const { setCellPopoverProps, DefaultCellPopover } = props;
 
     useEffect(() => {
@@ -30,4 +32,10 @@ export const getCustomCellPopoverRenderer = () => {
 
     return <DefaultCellPopover {...props} />;
   });
+
+  // Components passed to EUI DataGrid cannot be memoized components
+  // otherwise EUI throws an error `typeof Component !== 'function'`
+  return (props: EuiDataGridCellPopoverElementProps) => (
+    <RenderCustomCellPopoverMemoized {...props} />
+  );
 };

packages/kbn-unified-data-table/src/utils/get_render_cell_value.tsx

@@ -6,7 +6,7 @@
  * Side Public License, v 1.
  */
 
-import React, { memo, useEffect, useContext } from 'react';
+import React, { useEffect, useContext } from 'react';
 import { i18n } from '@kbn/i18n';
 import type { DataView, DataViewField } from '@kbn/data-views-plugin/public';
 import {
@@ -47,10 +47,7 @@ export const getRenderCellValueFn = ({
   externalCustomRenderers?: CustomCellRenderer;
   isPlainRecord?: boolean;
 }) => {
-  /**
-   * memo is imperative here otherwise the cell will re-render on every hover on every cell
-   */
-  return memo(function UnifiedDataTableRenderCellValue({
+  return function UnifiedDataTableRenderCellValue({
     rowIndex,
     columnId,
     isDetails,
@@ -149,7 +146,7 @@ export const getRenderCellValueFn = ({
         }}
       />
     );
-  });
+  };
 };
 
 /**

x-pack/plugins/security_solution/public/common/components/header_actions/actions.tsx

@@ -11,14 +11,15 @@ import { EuiButtonIcon, EuiToolTip } from '@elastic/eui';
 import styled from 'styled-components';
 
 import { TimelineTabs, TableId } from '@kbn/securitysolution-data-table';
+import { selectTimelineById } from '../../../timelines/store/selectors';
 import {
   eventHasNotes,
   getEventType,
   getPinOnClick,
 } from '../../../timelines/components/timeline/body/helpers';
 import { getScopedActions, isTimelineScope } from '../../../helpers';
 import { useIsInvestigateInResolverActionEnabled } from '../../../detections/components/alerts_table/timeline_actions/investigate_in_resolver';
-import { timelineActions, timelineSelectors } from '../../../timelines/store';
+import { timelineActions } from '../../../timelines/store';
 import type { ActionProps, OnPinEvent } from '../../../../common/types';
 import { TimelineId } from '../../../../common/types';
 import { AddEventNoteAction } from './add_note_icon_item';
@@ -66,11 +67,10 @@ const ActionsComponent: React.FC<ActionProps> = ({
     'unifiedComponentsInTimelineEnabled'
   );
   const emptyNotes: string[] = [];
-  const getTimeline = useMemo(() => timelineSelectors.getTimelineByIdSelector(), []);
-  const timelineType = useShallowEqualSelector(
-    (state) =>
-      (isTimelineScope(timelineId) ? getTimeline(state, timelineId) : timelineDefaults).timelineType
+  const { timelineType } = useShallowEqualSelector((state) =>
+    isTimelineScope(timelineId) ? selectTimelineById(state, timelineId) : timelineDefaults
   );
+
   const { startTransaction } = useStartTransaction();
 
   const isEnterprisePlus = useLicense().isEnterprise();
@@ -213,8 +213,8 @@ const ActionsComponent: React.FC<ActionProps> = ({
     onEventDetailsPanelOpened();
   }, [activeStep, incrementStep, isTourAnchor, isTourShown, onEventDetailsPanelOpened]);
   const showExpandEvent = useMemo(
-    () => !unifiedComponentsInTimelineEnabled || isEventViewer || timelineId !== TimelineId.active,
-    [isEventViewer, timelineId, unifiedComponentsInTimelineEnabled]
+    () => !unifiedComponentsInTimelineEnabled || isEventViewer,
+    [isEventViewer, unifiedComponentsInTimelineEnabled]
   );
 
   return (

x-pack/plugins/security_solution/public/common/components/header_actions/header_actions.tsx

@@ -16,8 +16,7 @@ import { TimelineTabs, TimelineId } from '../../../../common/types';
 import { isFullScreen } from '../../../timelines/components/timeline/body/column_headers';
 import { isActiveTimeline } from '../../../helpers';
 import { getColumnHeader } from '../../../timelines/components/timeline/body/column_headers/helpers';
-import { timelineActions, timelineSelectors } from '../../../timelines/store';
-import { useDeepEqualSelector } from '../../hooks/use_selector';
+import { timelineActions } from '../../../timelines/store';
 import { useGlobalFullScreen, useTimelineFullScreen } from '../../containers/use_full_screen';
 import { useKibana } from '../../lib/kibana';
 import { DEFAULT_ACTION_BUTTON_WIDTH } from '.';
@@ -27,6 +26,8 @@ import { EXIT_FULL_SCREEN } from '../exit_full_screen/translations';
 import { EventsSelect } from '../../../timelines/components/timeline/body/column_headers/events_select';
 import * as i18n from './translations';
 import { useIsExperimentalFeatureEnabled } from '../../hooks/use_experimental_features';
+import { useDeepEqualSelector } from '../../hooks/use_selector';
+import { selectTimelineById } from '../../../timelines/store/selectors';
 
 const SortingColumnsContainer = styled.div`
   button {
@@ -90,14 +91,14 @@ const HeaderActionsComponent: React.FC<HeaderActionProps> = memo(
     const { timelineFullScreen, setTimelineFullScreen } = useTimelineFullScreen();
     const dispatch = useDispatch();
 
-    const getManageTimeline = useMemo(() => timelineSelectors.getTimelineByIdSelector(), []);
-    const { defaultColumns } = useDeepEqualSelector((state) =>
-      getManageTimeline(state, timelineId)
-    );
     const unifiedComponentsInTimelineEnabled = useIsExperimentalFeatureEnabled(
       'unifiedComponentsInTimelineEnabled'
     );
 
+    const { defaultColumns } = useDeepEqualSelector((state) =>
+      selectTimelineById(state, timelineId)
+    );
+
     const toggleFullScreen = useCallback(() => {
       if (timelineId === TimelineId.active) {
         setTimelineFullScreen(!timelineFullScreen);

x-pack/plugins/security_solution/public/detections/components/alerts_table/timeline_actions/alert_context_menu.tsx

@@ -152,8 +152,9 @@ const AlertContextMenuComponent: React.FC<AlertContextMenuProps> = ({
       refetchQuery([timelineQuery]);
     } else {
       refetchQuery(globalQuery);
-      if (refetch) refetch();
     }
+
+    if (refetch) refetch();
   }, [scopeId, globalQuery, timelineQuery, refetch]);
 
   const ruleIndex =

x-pack/plugins/security_solution/public/timelines/components/timeline/body/events/event_column_view.test.tsx

@@ -145,14 +145,16 @@ describe('EventColumnView', () => {
   });
 
   test('it renders correct tooltip for NotesButton - timeline template', () => {
-    (useShallowEqualSelector as jest.Mock).mockReturnValue(TimelineType.template);
+    (useShallowEqualSelector as jest.Mock).mockReturnValue({
+      timelineType: TimelineType.template,
+    });
 
     const wrapper = mount(<EventColumnView {...props} />, { wrappingComponent: TestProviders });
 
     expect(wrapper.find('[data-test-subj="add-note"]').prop('toolTip')).toEqual(
       NOTES_DISABLE_TOOLTIP
     );
-    (useShallowEqualSelector as jest.Mock).mockReturnValue(TimelineType.default);
+    (useShallowEqualSelector as jest.Mock).mockReturnValue({ timelineType: TimelineType.default });
   });
 
   test('it does NOT render a pin button when isEventViewer is true', () => {

x-pack/plugins/security_solution/public/timelines/components/timeline/body/unified_timeline_body.test.tsx

@@ -46,6 +46,8 @@ const defaultProps: UnifiedTimelineBodyProps = {
     activePage: 0,
     querySize: 0,
   },
+  eventIdToNoteIds: {} as Record<string, string[]>,
+  pinnedEventIds: {} as Record<string, boolean>,
 };
 
 const renderTestComponents = (props?: UnifiedTimelineBodyProps) => {

x-pack/plugins/security_solution/public/timelines/components/timeline/tabs/eql/index.tsx

@@ -92,8 +92,6 @@ export const EqlTabContentComponent: React.FC<Props> = ({
   } = useSourcererDataView(SourcererScopeName.timeline);
   const { augmentedColumnHeaders, timelineQueryFieldsFromColumns } = useTimelineColumns(columns);
 
-  const leadingControlColumns = useTimelineControlColumn(columns, TIMELINE_NO_SORTING);
-
   const unifiedComponentsInTimelineEnabled = useIsExperimentalFeatureEnabled(
     'unifiedComponentsInTimelineEnabled'
   );
@@ -137,6 +135,14 @@ export const EqlTabContentComponent: React.FC<Props> = ({
     timerangeKind,
   });
 
+  const leadingControlColumns = useTimelineControlColumn({
+    columns,
+    sort: TIMELINE_NO_SORTING,
+    timelineId,
+    activeTab: TimelineTabs.eql,
+    refetch,
+  });
+
   const isQueryLoading = useMemo(
     () =>
       dataLoadingState === DataLoadingState.loading ||

x-pack/plugins/security_solution/public/timelines/components/timeline/tabs/pinned/index.tsx

@@ -171,7 +171,13 @@ export const PinnedTabContentComponent: React.FC<Props> = ({
       timerangeKind: undefined,
     });
 
-  const leadingControlColumns = useTimelineControlColumn(columns, sort);
+  const leadingControlColumns = useTimelineControlColumn({
+    columns,
+    sort,
+    timelineId,
+    activeTab: TimelineTabs.pinned,
+    refetch,
+  });
 
   const isQueryLoading = useMemo(
     () => [DataLoadingState.loading, DataLoadingState.loadingMore].includes(queryLoadingState),

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

@@ -203,7 +203,13 @@ export const QueryTabContentComponent: React.FC<Props> = ({
     timerangeKind,
   });
 
-  const leadingControlColumns = useTimelineControlColumn(columns, sort);
+  const leadingControlColumns = useTimelineControlColumn({
+    columns,
+    sort,
+    timelineId,
+    activeTab: TimelineTabs.query,
+    refetch,
+  });
 
   useEffect(() => {
     dispatch(