From 070b19f79dd3ad3d317fdaf6939d91e641e53c15 Mon Sep 17 00:00:00 2001
From: Maxim Palenov <maxim.palenov@elastic.co>
Date: Thu, 14 Nov 2024 11:59:45 +0100
Subject: [PATCH] add persistence query tests

---
 .../select_rule_type/test_helpers.ts          |  29 ++
 .../step_define_rule/index.test.tsx           | 383 +++++++++++++-----
 .../components/step_define_rule/index.tsx     |  75 ++--
 .../step_define_rule/use_persistent_query.ts  |  57 ++-
 .../pages/rule_creation/index.tsx             |   1 -
 .../pages/rule_editing/index.tsx              |   1 -
 6 files changed, 405 insertions(+), 141 deletions(-)
 create mode 100644 x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/components/select_rule_type/test_helpers.ts

diff --git a/x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/components/select_rule_type/test_helpers.ts b/x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/components/select_rule_type/test_helpers.ts
new file mode 100644
index 000000000000..a4d5fca3a2ff
--- /dev/null
+++ b/x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/components/select_rule_type/test_helpers.ts
@@ -0,0 +1,29 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+// eslint-disable-next-line import/no-extraneous-dependencies
+import { act, fireEvent, within, screen } from '@testing-library/react';
+import type { Type as RuleType } from '@kbn/securitysolution-io-ts-alerting-types';
+
+export async function selectRuleType(ruleType: RuleType): Promise<void> {
+  const testId = RULE_TYPE_TEST_ID_MAP[ruleType];
+
+  await act(async () => fireEvent.click(screen.getByTestId(testId)));
+
+  expect(within(screen.getByTestId(testId)).getByRole('switch')).toBeChecked();
+}
+
+const RULE_TYPE_TEST_ID_MAP = {
+  query: 'customRuleType',
+  saved_query: 'customRuleType',
+  eql: 'eqlRuleType',
+  machine_learning: 'machineLearningRuleType',
+  threshold: 'thresholdRuleType',
+  threat_match: 'threatMatchRuleType',
+  new_terms: 'newTermsRuleType',
+  esql: 'esqlRuleType',
+};
diff --git a/x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/components/step_define_rule/index.test.tsx b/x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/components/step_define_rule/index.test.tsx
index 50264fffabfb..b38c7b988419 100644
--- a/x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/components/step_define_rule/index.test.tsx
+++ b/x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/components/step_define_rule/index.test.tsx
@@ -5,6 +5,7 @@
  * 2.0.
  */
 
+import type { ChangeEvent } from 'react';
 import React, { useEffect, useState } from 'react';
 import { screen, fireEvent, render, within, act, waitFor } from '@testing-library/react';
 import type { Type as RuleType } from '@kbn/securitysolution-io-ts-alerting-types';
@@ -42,13 +43,43 @@ import {
   addRelatedIntegrationRow,
   setVersion,
 } from '../../../rule_creation/components/related_integrations/test_helpers';
+import { useEsqlAvailability } from '../../../../common/hooks/esql/use_esql_availability';
+import { useMLRuleConfig } from '../../../../common/components/ml/hooks/use_ml_rule_config';
+import { selectRuleType } from '../select_rule_type/test_helpers';
+
+// Set the extended default timeout for all define rule step form test
+jest.setTimeout(10 * 1000);
 
 // Mocks integrations
 jest.mock('../../../fleet_integrations/api');
+
+const MOCKED_QUERY_BAR_TEST_ID = 'mockedQueryBar';
+const MOCKED_LANGUAGE_INPUT_TEST_ID = 'languageInput';
+
+// Mocking QueryBar to avoid pulling and mocking a ton of dependencies
 jest.mock('../../../../common/components/query_bar', () => {
   return {
-    QueryBar: jest.fn(({ filterQuery }) => {
-      return <div data-test-subj="query-bar">{`${filterQuery.query} ${filterQuery.language}`}</div>;
+    QueryBar: jest.fn().mockImplementation(({ filterQuery, onSubmitQuery }) => {
+      const handleQueryChange = (event: ChangeEvent<HTMLTextAreaElement>) => {
+        onSubmitQuery({ query: event.target.value, language: filterQuery.language });
+      };
+
+      const handleLanguageChange = (event: ChangeEvent<HTMLInputElement>) => {
+        onSubmitQuery({ query: filterQuery.query, language: event.target.value });
+      };
+
+      return (
+        <div data-test-subj={MOCKED_QUERY_BAR_TEST_ID}>
+          <textarea value={filterQuery.query} onChange={handleQueryChange} />
+          <input
+            role="listbox"
+            type="text"
+            value={filterQuery.language}
+            onChange={handleLanguageChange}
+            data-test-subj={MOCKED_LANGUAGE_INPUT_TEST_ID}
+          />
+        </div>
+      );
     }),
   };
 });
@@ -77,6 +108,7 @@ const mockRedirectLegacyUrl = jest.fn();
 const mockGetLegacyUrlConflict = jest.fn();
 jest.mock('../../../../common/lib/kibana', () => {
   const originalModule = jest.requireActual('../../../../common/lib/kibana');
+
   return {
     ...originalModule,
     useToasts: jest.fn().mockReturnValue({
@@ -85,41 +117,43 @@ jest.mock('../../../../common/lib/kibana', () => {
       addWarning: jest.fn(),
       remove: jest.fn(),
     }),
-    useKibana: () => ({
-      services: {
-        ...originalModule.useKibana().services,
-        storage: {
-          get: jest.fn().mockReturnValue(true),
-        },
-        application: {
-          getUrlForApp: (appId: string, options?: { path?: string }) =>
-            `/app/${appId}${options?.path}`,
-          navigateToApp: jest.fn(),
-          capabilities: {
-            actions: {
-              delete: true,
-              save: true,
-              show: true,
+    useKibana: () => {
+      return {
+        services: {
+          ...originalModule.useKibana().services,
+          storage: {
+            get: jest.fn().mockReturnValue(true),
+          },
+          application: {
+            getUrlForApp: (appId: string, options?: { path?: string }) =>
+              `/app/${appId}${options?.path}`,
+            navigateToApp: jest.fn(),
+            capabilities: {
+              actions: {
+                delete: true,
+                save: true,
+                show: true,
+              },
             },
           },
-        },
-        data: {
-          search: {
-            search: () => ({
-              subscribe: () => ({
-                unsubscribe: jest.fn(),
+          data: {
+            search: {
+              search: () => ({
+                subscribe: () => ({
+                  unsubscribe: jest.fn(),
+                }),
               }),
-            }),
+            },
           },
-        },
-        spaces: {
-          ui: {
-            components: { getLegacyUrlConflict: mockGetLegacyUrlConflict },
-            redirectLegacyUrl: mockRedirectLegacyUrl,
+          spaces: {
+            ui: {
+              components: { getLegacyUrlConflict: mockGetLegacyUrlConflict },
+              redirectLegacyUrl: mockRedirectLegacyUrl,
+            },
           },
         },
-      },
-    }),
+      };
+    },
   };
 });
 jest.mock('../../../../common/hooks/use_selector', () => {
@@ -172,18 +206,28 @@ jest.mock('react-redux', () => {
 
 jest.mock('../../../../detections/containers/detection_engine/rules/use_rule_from_timeline');
 
+jest.mock('../../../../common/hooks/esql/use_esql_availability');
+jest.mock('../../../../common/components/ml/hooks/use_ml_rule_config');
+
 const mockUseRuleFromTimeline = useRuleFromTimeline as jest.Mock;
 const onOpenTimeline = jest.fn();
 
 const COMBO_BOX_TOGGLE_BUTTON_TEST_ID = 'comboBoxToggleListButton';
 const VERSION_INPUT_TEST_ID = 'relatedIntegrationVersionDependency';
+const DEFINE_RULE_FORM_STEP = 'defineRuleFormStepQueryEditor';
 
-// Failing: See https://github.com/elastic/kibana/issues/199648
-// Failing: See https://github.com/elastic/kibana/issues/199700
-describe.skip('StepDefineRule', () => {
+describe('StepDefineRule', () => {
   beforeEach(() => {
-    jest.clearAllMocks();
+    // jest.clearAllMocks();
     mockUseRuleFromTimeline.mockReturnValue({ onOpenTimeline, loading: false });
+    (useEsqlAvailability as jest.Mock).mockReturnValue({ isEsqlRuleTypeEnabled: true });
+    (useMLRuleConfig as jest.Mock).mockReturnValue({
+      allJobsStarted: true,
+      hasMlAdminPermissions: true,
+      hasMlLicense: true,
+      loading: false,
+      mlSuppressionFields: [],
+    });
   });
 
   it('renders correctly', () => {
@@ -194,6 +238,151 @@ describe.skip('StepDefineRule', () => {
     expect(screen.getByTestId('stepDefineRule')).toBeDefined();
   });
 
+  describe('query', () => {
+    it.each([
+      ['query', 'eql'],
+      // ['query', 'esql'],
+    ] as RuleType[][])(
+      'persists kuery when switching between "%s" and "%s" rule types',
+      async (ruleTypeA, ruleTypeB) => {
+        render(<TestForm />, {
+          wrapper: TestProviders,
+        });
+
+        typeQuery('someField : *');
+        selectQueryLanguage('kuery');
+
+        expectQuery(DEFINE_RULE_FORM_STEP, 'someField : *');
+
+        await selectRuleType(ruleTypeB);
+
+        expectQuery(DEFINE_RULE_FORM_STEP, '');
+
+        await selectRuleType(ruleTypeA);
+
+        expectQuery(DEFINE_RULE_FORM_STEP, 'someField : *');
+      }
+    );
+
+    it.each([
+      ['query', 'threshold'],
+      ['query', 'new_terms'],
+    ] as RuleType[][])(
+      'persists kuery when switching between "%s" and "%s" rule types',
+      async (ruleTypeA, ruleTypeB) => {
+        render(<TestForm />, {
+          wrapper: TestProviders,
+        });
+
+        typeQuery('someField : *');
+        selectQueryLanguage('kuery');
+
+        expectQuery(DEFINE_RULE_FORM_STEP, 'someField : *');
+
+        await selectRuleType(ruleTypeB);
+
+        expectQuery(DEFINE_RULE_FORM_STEP, 'someField : *');
+
+        await selectRuleType(ruleTypeA);
+
+        expectQuery(DEFINE_RULE_FORM_STEP, 'someField : *');
+      }
+    );
+
+    it.each([['query'], ['threshold'], ['threat_match'], ['new_terms']] as RuleType[][])(
+      'persists kuery when switching between "%s" and "threat_match"  rule types',
+      async (ruleType) => {
+        render(<TestForm />, {
+          wrapper: TestProviders,
+        });
+        await selectRuleType(ruleType);
+        typeQuery('someField : *');
+
+        expectQuery(DEFINE_RULE_FORM_STEP, 'someField : *');
+
+        await selectRuleType('threat_match');
+
+        expectQuery(DEFINE_RULE_FORM_STEP, 'someField : *');
+
+        await selectRuleType(ruleType);
+
+        expectQuery(DEFINE_RULE_FORM_STEP, 'someField : *');
+      }
+    );
+
+    it.each([['query'], ['esql'], ['threshold'], ['new_terms']] as RuleType[][])(
+      'persists EQL query when switching between "eql" and "%s" rule types',
+      async (switchRuleType) => {
+        render(<TestForm />, {
+          wrapper: TestProviders,
+        });
+
+        await selectRuleType('eql');
+        typeQuery('process where process.name == "regsvr32.exe"');
+
+        expectQuery(DEFINE_RULE_FORM_STEP, 'process where process.name == "regsvr32.exe"');
+
+        await selectRuleType(switchRuleType);
+
+        expectQuery(DEFINE_RULE_FORM_STEP, '');
+
+        await selectRuleType('eql');
+
+        expectQuery(DEFINE_RULE_FORM_STEP, 'process where process.name == "regsvr32.exe"');
+      }
+    );
+
+    it.each([['query'], ['eql'], ['threshold'], ['new_terms']] as RuleType[][])(
+      'persists ES|QL query when switching between "esql" and "%s" rule types',
+      async (switchRuleType) => {
+        render(<TestForm />, {
+          wrapper: TestProviders,
+        });
+        await selectRuleType('esql');
+        typeQuery('from test*');
+
+        expectQuery(DEFINE_RULE_FORM_STEP, 'from test*');
+
+        await selectRuleType(switchRuleType);
+
+        expectQuery(DEFINE_RULE_FORM_STEP, '');
+
+        await selectRuleType('esql');
+
+        expectQuery(DEFINE_RULE_FORM_STEP, 'from test*');
+      }
+    );
+
+    it('makes sure "threat_match" rule has "*:*" kuery by default', async () => {
+      render(<TestForm />, {
+        wrapper: TestProviders,
+      });
+      await selectRuleType('threat_match');
+
+      expectQuery(DEFINE_RULE_FORM_STEP, '*:*');
+    });
+
+    it.each([['query'], ['eql'], ['esql'], ['threshold'], ['new_terms']] as RuleType[][])(
+      'makes sure "threat_match" rule has "*:*" kuery when switched from "%s" rule type without user input',
+      async (ruleType) => {
+        render(<TestForm />, {
+          wrapper: TestProviders,
+        });
+        await selectRuleType(ruleType);
+
+        expectQuery(DEFINE_RULE_FORM_STEP, '');
+
+        await selectRuleType('threat_match');
+
+        expectQuery(DEFINE_RULE_FORM_STEP, '*:*');
+
+        await selectRuleType(ruleType);
+
+        expectQuery(DEFINE_RULE_FORM_STEP, '');
+      }
+    );
+  });
+
   describe('alert suppression', () => {
     it('persists state when switching between custom query and threshold rule types', async () => {
       const mockFields: FieldSpec[] = [
@@ -205,7 +394,7 @@ describe.skip('StepDefineRule', () => {
         },
       ];
 
-      const { rerender } = render(
+      render(
         <TestForm
           indexPattern={{
             title: '',
@@ -221,29 +410,11 @@ describe.skip('StepDefineRule', () => {
       setDurationType('Per time period');
       setDuration(10, 'h');
 
-      // switch to threshold rule type
-      rerender(
-        <TestForm
-          ruleType="threshold"
-          indexPattern={{
-            title: '',
-            fields: mockFields,
-          }}
-        />
-      );
+      await selectRuleType('threshold');
 
       expectDuration(10, 'h');
 
-      // switch back to custom query rule type
-      rerender(
-        <TestForm
-          ruleType="query"
-          indexPattern={{
-            title: '',
-            fields: mockFields,
-          }}
-        />
-      );
+      await selectRuleType('query');
 
       expectSuppressionFields(['test-field']);
       expectDuration(10, 'h');
@@ -479,7 +650,7 @@ describe.skip('StepDefineRule', () => {
 
   describe('handleSetRuleFromTimeline', () => {
     it('updates KQL query correctly', () => {
-      const kqlQuery = {
+      const timelineKqlQuery = {
         index: ['.alerts-security.alerts-default', 'logs-*', 'packetbeat-*'],
         queryBar: {
           filters: [],
@@ -493,7 +664,7 @@ describe.skip('StepDefineRule', () => {
 
       mockUseRuleFromTimeline.mockImplementation((handleSetRuleFromTimeline) => {
         useEffect(() => {
-          handleSetRuleFromTimeline(kqlQuery);
+          handleSetRuleFromTimeline(timelineKqlQuery);
         }, [handleSetRuleFromTimeline]);
 
         return { onOpenTimeline, loading: false };
@@ -503,13 +674,12 @@ describe.skip('StepDefineRule', () => {
         wrapper: TestProviders,
       });
 
-      expect(screen.getAllByTestId('query-bar')[0].textContent).toBe(
-        `${kqlQuery.queryBar.query.query} ${kqlQuery.queryBar.query.language}`
-      );
+      expectQuery(DEFINE_RULE_FORM_STEP, timelineKqlQuery.queryBar.query.query);
+      expectQueryLanguage(DEFINE_RULE_FORM_STEP, timelineKqlQuery.queryBar.query.language);
     });
 
     it('updates EQL query correctly', async () => {
-      const eqlQuery = {
+      const timelineEqlQuery = {
         index: ['.alerts-security.alerts-default', 'logs-*', 'packetbeat-*'],
         queryBar: {
           filters: [],
@@ -528,20 +698,26 @@ describe.skip('StepDefineRule', () => {
         },
       };
 
+      const setRuleFromTimelineFactory =
+        (handleSetRuleFromTimeline: (value: unknown) => void) => async () =>
+          act(async () => handleSetRuleFromTimeline(timelineEqlQuery));
+
+      let setRuleFromTimeline!: () => void;
+
       mockUseRuleFromTimeline.mockImplementation((handleSetRuleFromTimeline) => {
-        useEffect(() => {
-          handleSetRuleFromTimeline(eqlQuery);
-        }, [handleSetRuleFromTimeline]);
+        setRuleFromTimeline = setRuleFromTimelineFactory(handleSetRuleFromTimeline);
 
         return { onOpenTimeline, loading: false };
       });
 
-      render(<TestForm ruleType="eql" />, {
+      render(<TestForm />, {
         wrapper: TestProviders,
       });
 
-      expect(screen.getByTestId(`eqlQueryBarTextInput`).textContent).toEqual(
-        eqlQuery.queryBar.query.query
+      await setRuleFromTimeline();
+
+      expect(screen.getByTestId('eqlQueryBarTextInput')).toHaveValue(
+        timelineEqlQuery.queryBar.query.query
       );
 
       await act(async () => {
@@ -550,14 +726,14 @@ describe.skip('StepDefineRule', () => {
 
       expect(
         within(screen.getByTestId('eql-event-category-field')).queryByRole('combobox')
-      ).toHaveValue(eqlQuery.eqlOptions.eventCategoryField);
+      ).toHaveValue(timelineEqlQuery.eqlOptions.eventCategoryField);
 
       expect(
         within(screen.getByTestId('eql-tiebreaker-field')).queryByRole('combobox')
-      ).toHaveValue(eqlQuery.eqlOptions.tiebreakerField);
+      ).toHaveValue(timelineEqlQuery.eqlOptions.tiebreakerField);
 
       expect(within(screen.getByTestId('eql-timestamp-field')).queryByRole('combobox')).toHaveValue(
-        eqlQuery.eqlOptions.timestampField
+        timelineEqlQuery.eqlOptions.timestampField
       );
     });
   });
@@ -571,15 +747,9 @@ describe.skip('StepDefineRule', () => {
           saved_id: null,
         },
       };
-      render(
-        <TestForm
-          formProps={{ isQueryBarValid: false, ruleType: 'query' }}
-          initialState={initialState}
-        />,
-        {
-          wrapper: TestProviders,
-        }
-      );
+      render(<TestForm formProps={{ isQueryBarValid: false }} initialState={initialState} />, {
+        wrapper: TestProviders,
+      });
 
       expect(screen.getByTestId('ai-assistant')).toBeInTheDocument();
     });
@@ -592,31 +762,26 @@ describe.skip('StepDefineRule', () => {
           saved_id: null,
         },
       };
-      render(
-        <TestForm
-          formProps={{ isQueryBarValid: false, ruleType: 'query' }}
-          initialState={initialState}
-        />,
-        {
-          wrapper: TestProviders,
-        }
-      );
+      render(<TestForm formProps={{ isQueryBarValid: false }} initialState={initialState} />, {
+        wrapper: TestProviders,
+      });
 
       expect(screen.queryByTestId('ai-assistant')).toBe(null);
     });
 
     it('does not render assistant when query is valid', () => {
-      render(<TestForm formProps={{ isQueryBarValid: true, ruleType: 'query' }} />, {
+      render(<TestForm formProps={{ isQueryBarValid: true }} />, {
         wrapper: TestProviders,
       });
 
       expect(screen.queryByTestId('ai-assistant')).toBe(null);
     });
 
-    it('does not render assistant for ML rule type', () => {
-      render(<TestForm formProps={{ isQueryBarValid: false, ruleType: 'machine_learning' }} />, {
+    it('does not render assistant for ML rule type', async () => {
+      render(<TestForm formProps={{ isQueryBarValid: false }} />, {
         wrapper: TestProviders,
       });
+      await selectRuleType('machine_learning');
 
       expect(screen.queryByTestId('ai-assistant')).toBe(null);
     });
@@ -625,7 +790,6 @@ describe.skip('StepDefineRule', () => {
 
 interface TestFormProps {
   initialState?: Partial<DefineStepRule>;
-  ruleType?: RuleType;
   indexPattern?: DataViewBase;
   onSubmit?: FormSubmitHandler<DefineStepRule>;
   formProps?: Partial<StepDefineRuleProps>;
@@ -633,7 +797,6 @@ interface TestFormProps {
 
 function TestForm({
   initialState,
-  ruleType = stepDefineDefaultValue.ruleType,
   indexPattern = { fields: [], title: '' },
   onSubmit,
   formProps,
@@ -660,7 +823,6 @@ function TestForm({
         isQueryBarValid={true}
         setIsQueryBarValid={jest.fn()}
         setIsThreatQueryBarValid={jest.fn()}
-        ruleType={ruleType}
         index={stepDefineDefaultValue.index}
         threatIndex={stepDefineDefaultValue.threatIndex}
         alertSuppressionFields={stepDefineDefaultValue[ALERT_SUPPRESSION_FIELDS_FIELD_NAME]}
@@ -683,3 +845,34 @@ function submitForm(): Promise<void> {
     fireEvent.click(screen.getByText('Submit'));
   });
 }
+
+function typeQuery(query: string): void {
+  act(() => {
+    fireEvent.input(within(screen.getByTestId(DEFINE_RULE_FORM_STEP)).getByRole('textbox'), {
+      target: { value: query },
+    });
+  });
+}
+
+function expectQuery(containerTestId: string, expectedQuery: string): void {
+  expect(within(screen.getByTestId(containerTestId)).getByRole('textbox')).toHaveValue(
+    expectedQuery
+  );
+}
+
+function selectQueryLanguage(language: 'kuery' | 'eql' | 'esql'): void {
+  act(() => {
+    fireEvent.input(
+      within(screen.getByTestId(DEFINE_RULE_FORM_STEP)).getByTestId(MOCKED_LANGUAGE_INPUT_TEST_ID),
+      {
+        target: { value: language },
+      }
+    );
+  });
+}
+
+function expectQueryLanguage(containerTestId: string, expectedLanguage: string): void {
+  expect(
+    within(screen.getByTestId(containerTestId)).getByTestId(MOCKED_LANGUAGE_INPUT_TEST_ID)
+  ).toHaveValue(expectedLanguage);
+}
diff --git a/x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/components/step_define_rule/index.tsx b/x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/components/step_define_rule/index.tsx
index 0abe888e5148..62dc978c01bf 100644
--- a/x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/components/step_define_rule/index.tsx
+++ b/x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/components/step_define_rule/index.tsx
@@ -22,7 +22,6 @@ import styled from 'styled-components';
 import { i18n as i18nCore } from '@kbn/i18n';
 import { isEqual, isEmpty } from 'lodash';
 import type { FieldSpec } from '@kbn/data-plugin/common';
-import type { Type } from '@kbn/securitysolution-io-ts-alerting-types';
 
 import type { SavedQuery } from '@kbn/data-plugin/public';
 import type { DataViewBase } from '@kbn/es-query';
@@ -113,7 +112,6 @@ export interface StepDefineRuleProps extends RuleStepProps {
   isQueryBarValid: boolean;
   setIsQueryBarValid: (valid: boolean) => void;
   setIsThreatQueryBarValid: (valid: boolean) => void;
-  ruleType: Type;
   index: string[];
   threatIndex: string[];
   alertSuppressionFields?: string[];
@@ -167,7 +165,6 @@ const StepDefineRuleComponent: FC<StepDefineRuleProps> = ({
   optionsSelected,
   queryBarSavedId,
   queryBarTitle,
-  ruleType,
   setIsQueryBarValid,
   setIsThreatQueryBarValid,
   setOptionsSelected,
@@ -176,6 +173,11 @@ const StepDefineRuleComponent: FC<StepDefineRuleProps> = ({
   threatIndicesConfig,
   thresholdFields,
 }) => {
+  const [{ ruleType, queryBar: a }] = useFormData<DefineStepRule>({
+    form,
+    watch: ['ruleType', 'queryBar'],
+  });
+
   const { isSuppressionEnabled: isAlertSuppressionEnabled } = useAlertSuppression(ruleType);
   const [openTimelineSearch, setOpenTimelineSearch] = useState(false);
   const [indexModified, setIndexModified] = useState(false);
@@ -209,28 +211,6 @@ const StepDefineRuleComponent: FC<StepDefineRuleProps> = ({
     setFieldValue,
   });
 
-  const handleSetRuleFromTimeline = useCallback<SetRuleQuery>(
-    ({ index: timelineIndex, queryBar: timelineQueryBar, eqlOptions }) => {
-      const setQuery = () => {
-        setFieldValue('index', timelineIndex);
-        setFieldValue('queryBar', timelineQueryBar);
-      };
-      if (timelineQueryBar.query.language === 'eql') {
-        setRuleTypeCallback('eql', setQuery);
-        setOptionsSelected((prevOptions) => ({
-          ...prevOptions,
-          ...(eqlOptions != null ? eqlOptions : {}),
-        }));
-      } else {
-        setQuery();
-      }
-    },
-    [setFieldValue, setRuleTypeCallback, setOptionsSelected]
-  );
-
-  const { onOpenTimeline, loading: timelineQueryLoading } =
-    useRuleFromTimeline(handleSetRuleFromTimeline);
-
   // Callback for when user toggles between Data Views and Index Patterns
   const onChangeDataSource = useCallback(
     (optionId: string) => {
@@ -279,9 +259,36 @@ const StepDefineRuleComponent: FC<StepDefineRuleProps> = ({
     setThreatIndexModified(!isEqual(threatIndex, threatIndicesConfig));
   }, [threatIndex, threatIndicesConfig]);
 
-  usePersistentQuery({ form, ruleTypePath: 'ruleType', queryPath: 'queryBar' });
+  const { setPersistentEqlQuery } = usePersistentQuery({
+    form,
+    ruleTypePath: 'ruleType',
+    queryPath: 'queryBar',
+  });
   usePersistentAlertSuppressionState({ form });
 
+  const handleSetRuleFromTimeline = useCallback<SetRuleQuery>(
+    ({ index: timelineIndex, queryBar: timelineQueryBar, eqlOptions }) => {
+      const setQuery = () => {
+        setFieldValue('index', timelineIndex);
+        setFieldValue('queryBar', timelineQueryBar);
+      };
+      if (timelineQueryBar.query.language === 'eql') {
+        setPersistentEqlQuery(timelineQueryBar);
+        setRuleTypeCallback('eql', setQuery);
+        setOptionsSelected((prevOptions) => ({
+          ...prevOptions,
+          ...(eqlOptions != null ? eqlOptions : {}),
+        }));
+      } else {
+        setQuery();
+      }
+    },
+    [setFieldValue, setRuleTypeCallback, setOptionsSelected, setPersistentEqlQuery]
+  );
+
+  const { onOpenTimeline, loading: timelineQueryLoading } =
+    useRuleFromTimeline(handleSetRuleFromTimeline);
+
   // if saved query failed to load:
   // - reset shouldLoadFormDynamically to false, as non existent query cannot be used for loading and execution
   const handleSavedQueryError = useCallback(() => {
@@ -651,12 +658,18 @@ const StepDefineRuleComponent: FC<StepDefineRuleProps> = ({
             component={SelectRuleType}
             componentProps={selectRuleTypeProps}
           />
-          <RuleTypeEuiFormRow $isVisible={!isMlRule(ruleType)} fullWidth>
+          <RuleTypeEuiFormRow $isVisible={!isMlRule(ruleType) && !isEsqlRule(ruleType)} fullWidth>
+            <>
+              <EuiSpacer size="s" />
+              {DataSource}
+            </>
+          </RuleTypeEuiFormRow>
+          <RuleTypeEuiFormRow
+            $isVisible={!isMlRule(ruleType)}
+            fullWidth
+            data-test-subj="defineRuleFormStepQueryEditor"
+          >
             <>
-              <StyledVisibleContainer isVisible={!isEsqlRule(ruleType)}>
-                <EuiSpacer size="s" />
-                {DataSource}
-              </StyledVisibleContainer>
               <EuiSpacer size="s" />
               {isEqlRule(ruleType) ? (
                 <>
diff --git a/x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/components/step_define_rule/use_persistent_query.ts b/x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/components/step_define_rule/use_persistent_query.ts
index 5a79a580a7e1..d6e3cfd9605f 100644
--- a/x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/components/step_define_rule/use_persistent_query.ts
+++ b/x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/components/step_define_rule/use_persistent_query.ts
@@ -5,7 +5,7 @@
  * 2.0.
  */
 
-import { useEffect, useRef } from 'react';
+import { useEffect, useMemo, useRef } from 'react';
 import { isEqual } from 'lodash';
 import usePrevious from 'react-use/lib/usePrevious';
 import type { FieldHook } from '../../../../shared_imports';
@@ -33,14 +33,25 @@ interface UsePersistentQueryParams {
   queryPath: string;
 }
 
+interface UsePersistentQueryResult {
+  setPersistentQuery: (value: FieldValueQueryBar) => void;
+  setPersistentEqlQuery: (value: FieldValueQueryBar) => void;
+  setPersistentEsqlQuery: (value: FieldValueQueryBar) => void;
+}
+
 /**
- * Persists query when switching between different rule types using different queries (kuery, EQL, ES|QL).
+ * Persists query when switching between different rule types using different queries (kuery and ES|QL).
+ *
+ * When the user changes rule type to or from "threat_match" this will modify the
+ * default "Custom query" string to either:
+ *   * from '' to '*:*' if the type is switched to "threat_match"
+ *   * from '*:*' back to '' if the type is switched back from "threat_match" to another one
  */
 export function usePersistentQuery({
   form,
   ruleTypePath,
   queryPath,
-}: UsePersistentQueryParams): void {
+}: UsePersistentQueryParams): UsePersistentQueryResult {
   const [{ [ruleTypePath]: ruleType, [queryPath]: currentQuery }] = useFormData({
     form,
     watch: [ruleTypePath, queryPath],
@@ -55,13 +66,25 @@ export function usePersistentQuery({
       return;
     }
 
-    if (currentQuery?.query?.language === EQL_QUERY_LANGUAGE) {
-      eqlQueryRef.current = currentQuery;
-    } else if (currentQuery?.query?.language === ESQL_QUERY_LANGIAGE) {
-      esqlQueryRef.current = currentQuery;
-    } else {
-      queryRef.current = currentQuery;
+    if (isEqlRule(ruleType)) {
+      eqlQueryRef.current =
+        currentQuery?.query?.language === EQL_QUERY_LANGUAGE ? currentQuery : eqlQueryRef.current;
+
+      return;
+    }
+
+    if (isEsqlRule(ruleType)) {
+      esqlQueryRef.current =
+        currentQuery?.query?.language === ESQL_QUERY_LANGIAGE ? currentQuery : esqlQueryRef.current;
+
+      return;
+    }
+
+    if ([EQL_QUERY_LANGUAGE, ESQL_QUERY_LANGIAGE].includes(currentQuery?.query?.language)) {
+      return;
     }
+
+    queryRef.current = currentQuery;
   }, [ruleType, currentQuery]);
 
   useEffect(() => {
@@ -71,15 +94,14 @@ export function usePersistentQuery({
 
     const queryField = form.getFields()[queryPath] as FieldHook<FieldValueQueryBar>;
 
-    if (isEqlRule(ruleType) && queryField.value?.query?.language !== EQL_QUERY_LANGUAGE) {
+    if (isEqlRule(ruleType)) {
       queryField.reset({
         defaultValue: eqlQueryRef.current,
       });
-
       return;
     }
 
-    if (isEsqlRule(ruleType) && queryField.value?.query?.language !== ESQL_QUERY_LANGIAGE) {
+    if (isEsqlRule(ruleType)) {
       queryField.reset({
         defaultValue: esqlQueryRef.current,
       });
@@ -106,8 +128,17 @@ export function usePersistentQuery({
 
     if (isEqlRule(previousRuleType) || isEsqlRule(previousRuleType)) {
       queryField.reset({
-        defaultValue: queryRef.current ?? DEFAULT_KQL_QUERY_FIELD_VALUE,
+        defaultValue: queryRef.current,
       });
     }
   }, [queryPath, ruleType, previousRuleType, form]);
+
+  return useMemo(
+    () => ({
+      setPersistentQuery: (value: FieldValueQueryBar) => (queryRef.current = value),
+      setPersistentEqlQuery: (value: FieldValueQueryBar) => (eqlQueryRef.current = value),
+      setPersistentEsqlQuery: (value: FieldValueQueryBar) => (esqlQueryRef.current = value),
+    }),
+    [queryRef]
+  );
 }
diff --git a/x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/pages/rule_creation/index.tsx b/x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/pages/rule_creation/index.tsx
index d19c9c7c89d0..1aec784aef3c 100644
--- a/x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/pages/rule_creation/index.tsx
+++ b/x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/pages/rule_creation/index.tsx
@@ -563,7 +563,6 @@ const CreateRulePageComponent: React.FC = () => {
             isQueryBarValid={isQueryBarValid}
             setIsQueryBarValid={setIsQueryBarValid}
             setIsThreatQueryBarValid={setIsThreatQueryBarValid}
-            ruleType={defineStepData.ruleType}
             index={memoizedIndex}
             threatIndex={defineStepData.threatIndex}
             alertSuppressionFields={defineStepData[ALERT_SUPPRESSION_FIELDS_FIELD_NAME]}
diff --git a/x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/pages/rule_editing/index.tsx b/x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/pages/rule_editing/index.tsx
index 3b1112d47c91..77cb47eaf11a 100644
--- a/x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/pages/rule_editing/index.tsx
+++ b/x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/pages/rule_editing/index.tsx
@@ -236,7 +236,6 @@ const EditRulePageComponent: FC<{ rule: RuleResponse }> = ({ rule }) => {
                   isQueryBarValid={isQueryBarValid}
                   setIsQueryBarValid={setIsQueryBarValid}
                   setIsThreatQueryBarValid={setIsThreatQueryBarValid}
-                  ruleType={defineStepData.ruleType}
                   index={memoizedIndex}
                   threatIndex={defineStepData.threatIndex}
                   alertSuppressionFields={defineStepData[ALERT_SUPPRESSION_FIELDS_FIELD_NAME]}