Refactor useEnableDataFeed to split disablement and enablement logic

elastic · Apr 25, 2023 · 9981451 · 9981451
1 parent 34ebddc
commit 9981451
Show file tree

Hide file tree

Showing 8 changed files with 288 additions and 269 deletions.
diff --git a/...security_solution/public/common/components/ml_popover/hooks/use_enable_data_feed.test.tsx b/...security_solution/public/common/components/ml_popover/hooks/use_enable_data_feed.test.tsx
diff --git a/...ugins/security_solution/public/common/components/ml_popover/hooks/use_enable_data_feed.ts b/...ugins/security_solution/public/common/components/ml_popover/hooks/use_enable_data_feed.ts
@@ -20,17 +20,20 @@ import { setupMlJob, startDatafeeds, stopDatafeeds } from '../api';
 import type { ErrorResponse, SecurityJob } from '../types';
 import * as i18n from './translations';
 
-// Enable/Disable Job & Datafeed -- passed to JobsTable for use as callback on JobSwitch
+// Enable/Disable Job & Datafeed
 export const useEnableDataFeed = () => {
   const { telemetry } = useKibana().services;
 
   const { addError } = useAppToasts();
   const [isLoading, setIsLoading] = useState(false);
 
   const enableDatafeed = useCallback(
-    async (job: SecurityJob, latestTimestampMs: number, enable: boolean) => {
-      submitTelemetry(job, enable);
+    async (job: SecurityJob, latestTimestampMs: number) => {
       setIsLoading(true);
+      track(
+        METRIC_TYPE.COUNT,
+        job.isElasticJob ? TELEMETRY_EVENT.SIEM_JOB_ENABLED : TELEMETRY_EVENT.CUSTOM_JOB_ENABLED
+      );
 
       if (!job.isInstalled) {
         try {
@@ -40,16 +43,15 @@ export const useEnableDataFeed = () => {
             jobIdErrorFilter: [job.id],
             groups: job.groups,
           });
-          setIsLoading(false);
           telemetry.reportMLJobUpdate({
             jobId: job.id,
             isElasticJob: job.isElasticJob,
             moduleId: job.moduleId,
             status: ML_JOB_TELEMETRY_STATUS.moduleInstalled,
           });
         } catch (error) {
-          addError(error, { title: i18n.CREATE_JOB_FAILURE });
           setIsLoading(false);
+          addError(error, { title: i18n.CREATE_JOB_FAILURE });
           telemetry.reportMLJobUpdate({
             jobId: job.id,
             isElasticJob: job.isElasticJob,
@@ -67,94 +69,88 @@ export const useEnableDataFeed = () => {
       const maxStartTime = date.setDate(date.getDate() - 14);
 
       const datafeedId = `datafeed-${job.id}`;
-      if (enable) {
-        const startTime = Math.max(latestTimestampMs, maxStartTime);
-        const reportEnableJobError = (error: Error) => {
-          track(METRIC_TYPE.COUNT, TELEMETRY_EVENT.JOB_ENABLE_FAILURE);
-          addError(error, { title: i18n.START_JOB_FAILURE });
-          telemetry.reportMLJobUpdate({
-            jobId: job.id,
-            isElasticJob: job.isElasticJob,
-            status: ML_JOB_TELEMETRY_STATUS.startError,
-            errorMessage: `${i18n.START_JOB_FAILURE} - ${error.message}`,
-          });
-        };
-
-        try {
-          const response = await startDatafeeds({
-            datafeedIds: [datafeedId],
-            start: startTime,
-          });
 
-          if (response[datafeedId]?.error) {
-            throw new Error(response[datafeedId].error);
-          }
+      const startTime = Math.max(latestTimestampMs, maxStartTime);
 
-          telemetry.reportMLJobUpdate({
-            jobId: job.id,
-            isElasticJob: job.isElasticJob,
-            status: ML_JOB_TELEMETRY_STATUS.started,
-          });
+      try {
+        const response = await startDatafeeds({
+          datafeedIds: [datafeedId],
+          start: startTime,
+        });
 
-          return { enabled: response[datafeedId] ? response[datafeedId].started : false };
-        } catch (error) {
-          reportEnableJobError(error);
-        } finally {
-          setIsLoading(false);
+        if (response[datafeedId]?.error) {
+          throw new Error(response[datafeedId].error);
         }
-      } else {
-        const reportDisableError = (error: Error) => {
-          track(METRIC_TYPE.COUNT, TELEMETRY_EVENT.JOB_DISABLE_FAILURE);
-          addError(error, { title: i18n.STOP_JOB_FAILURE });
-          telemetry.reportMLJobUpdate({
-            jobId: job.id,
-            isElasticJob: job.isElasticJob,
-            status: ML_JOB_TELEMETRY_STATUS.stopError,
-            errorMessage: `${i18n.STOP_JOB_FAILURE} - ${error.message}`,
-          });
-        };
 
-        try {
-          const [response] = await stopDatafeeds({ datafeedIds: [datafeedId] });
+        telemetry.reportMLJobUpdate({
+          jobId: job.id,
+          isElasticJob: job.isElasticJob,
+          status: ML_JOB_TELEMETRY_STATUS.started,
+        });
+
+        return { enabled: response[datafeedId] ? response[datafeedId].started : false };
+      } catch (error) {
+        track(METRIC_TYPE.COUNT, TELEMETRY_EVENT.JOB_ENABLE_FAILURE);
+        addError(error, { title: i18n.START_JOB_FAILURE });
+        telemetry.reportMLJobUpdate({
+          jobId: job.id,
+          isElasticJob: job.isElasticJob,
+          status: ML_JOB_TELEMETRY_STATUS.startError,
+          errorMessage: `${i18n.START_JOB_FAILURE} - ${error.message}`,
+        });
+      } finally {
+        setIsLoading(false);
+      }
 
-          if (isErrorResponse(response)) {
-            throw new Error(response.error);
-          }
+      return { enabled: false };
+    },
+    [addError, telemetry]
+  );
 
-          telemetry.reportMLJobUpdate({
-            jobId: job.id,
-            isElasticJob: job.isElasticJob,
-            status: ML_JOB_TELEMETRY_STATUS.stopped,
-          });
+  const disableDatafeed = useCallback(
+    async (job: SecurityJob) => {
+      track(
+        METRIC_TYPE.COUNT,
+        job.isElasticJob ? TELEMETRY_EVENT.SIEM_JOB_DISABLED : TELEMETRY_EVENT.CUSTOM_JOB_DISABLED
+      );
+      setIsLoading(true);
 
-          return { enabled: response[datafeedId] ? !response[datafeedId].stopped : true };
-        } catch (error) {
-          reportDisableError(error);
-        } finally {
-          setIsLoading(false);
+      const datafeedId = `datafeed-${job.id}`;
+
+      try {
+        const [response] = await stopDatafeeds({ datafeedIds: [datafeedId] });
+
+        if (isErrorResponse(response)) {
+          throw new Error(response.error);
         }
+
+        telemetry.reportMLJobUpdate({
+          jobId: job.id,
+          isElasticJob: job.isElasticJob,
+          status: ML_JOB_TELEMETRY_STATUS.stopped,
+        });
+
+        return { enabled: response[datafeedId] ? !response[datafeedId].stopped : true };
+      } catch (error) {
+        track(METRIC_TYPE.COUNT, TELEMETRY_EVENT.JOB_DISABLE_FAILURE);
+        addError(error, { title: i18n.STOP_JOB_FAILURE });
+        telemetry.reportMLJobUpdate({
+          jobId: job.id,
+          isElasticJob: job.isElasticJob,
+          status: ML_JOB_TELEMETRY_STATUS.stopError,
+          errorMessage: `${i18n.STOP_JOB_FAILURE} - ${error.message}`,
+        });
+      } finally {
+        setIsLoading(false);
       }
-      return { enabled: !enable };
+
+      return { enabled: true };
     },
     [addError, telemetry]
   );
 
-  return { enableDatafeed, isLoading };
+  return { enableDatafeed, disableDatafeed, isLoading };
 };
 
 const isErrorResponse = (response: ErrorResponse): response is ErrorResponse =>
   !isEmpty(response.error);
-
-const submitTelemetry = (job: SecurityJob, enabled: boolean) => {
-  // Report type of job enabled/disabled
-  track(
-    METRIC_TYPE.COUNT,
-    job.isElasticJob
-      ? enabled
-        ? TELEMETRY_EVENT.SIEM_JOB_ENABLED
-        : TELEMETRY_EVENT.SIEM_JOB_DISABLED
-      : enabled
-      ? TELEMETRY_EVENT.CUSTOM_JOB_ENABLED
-      : TELEMETRY_EVENT.CUSTOM_JOB_DISABLED
-  );
-};
diff --git a/x-pack/plugins/security_solution/public/common/components/ml_popover/ml_popover.tsx b/x-pack/plugins/security_solution/public/common/components/ml_popover/ml_popover.tsx
@@ -59,13 +59,22 @@ export const MlPopover = React.memo(() => {
   } = useSecurityJobs();
 
   const docLinks = useKibana().services.docLinks;
-  const { enableDatafeed, isLoading: isLoadingEnableDataFeed } = useEnableDataFeed();
+  const {
+    enableDatafeed,
+    disableDatafeed,
+    isLoading: isLoadingEnableDataFeed,
+  } = useEnableDataFeed();
   const handleJobStateChange = useCallback(
     async (job: SecurityJob, latestTimestampMs: number, enable: boolean) => {
-      await enableDatafeed(job, latestTimestampMs, enable);
+      if (enable) {
+        await enableDatafeed(job, latestTimestampMs);
+      } else {
+        await disableDatafeed(job);
+      }
+
       refreshJobs();
     },
-    [refreshJobs, enableDatafeed]
+    [refreshJobs, enableDatafeed, disableDatafeed]
   );
 
   const filteredJobs = filterJobs({

diff --git a/...ins/security_solution/public/detection_engine/rule_management/logic/use_start_ml_jobs.tsx b/...ins/security_solution/public/detection_engine/rule_management/logic/use_start_ml_jobs.tsx
@@ -43,7 +43,7 @@ export const useStartMlJobs = (): ReturnUseStartMlJobs => {
           }
 
           const latestTimestampMs = job.latestTimestampMs ?? 0;
-          await enableDatafeed(job, latestTimestampMs, true);
+          await enableDatafeed(job, latestTimestampMs);
         })
       );
       refetchJobs();

diff --git a/.../components/rules/ml_jobs_description/admin/ml_admin_job_description.integration.test.tsx b/.../components/rules/ml_jobs_description/admin/ml_admin_job_description.integration.test.tsx
@@ -56,8 +56,7 @@ describe('MlAdminJobDescription', () => {
     userEvent.click(screen.getByTestId('job-switch'));
     expect(enableDatafeedSpy).toHaveBeenCalledWith(
       securityJobNotStarted,
-      securityJobNotStarted.latestTimestampMs,
-      true
+      securityJobNotStarted.latestTimestampMs
     );
 
     await waitFor(() => {

diff --git a/...public/detections/components/rules/ml_jobs_description/admin/ml_admin_job_description.tsx b/...public/detections/components/rules/ml_jobs_description/admin/ml_admin_job_description.tsx
@@ -25,14 +25,23 @@ const MlAdminJobDescriptionComponent: FC<MlAdminJobDescriptionProps> = ({
   loading,
   refreshJob,
 }) => {
-  const { enableDatafeed, isLoading: isLoadingEnableDataFeed } = useEnableDataFeed();
+  const {
+    enableDatafeed,
+    disableDatafeed,
+    isLoading: isLoadingEnableDataFeed,
+  } = useEnableDataFeed();
 
   const handleJobStateChange = useCallback(
     async (_, latestTimestampMs: number, enable: boolean) => {
-      await enableDatafeed(job, latestTimestampMs, enable);
+      if (enable) {
+        await enableDatafeed(job, latestTimestampMs);
+      } else {
+        await disableDatafeed(job);
+      }
+
       refreshJob(job);
     },
-    [enableDatafeed, job, refreshJob]
+    [enableDatafeed, disableDatafeed, job, refreshJob]
   );
 
   const switchComponent = useMemo(

diff --git a/...tion/public/overview/components/entity_analytics/anomalies/components/enable_job.test.tsx b/...tion/public/overview/components/entity_analytics/anomalies/components/enable_job.test.tsx
@@ -45,7 +45,7 @@ describe('EnableJob', () => {
     fireEvent.click(getByText('Run job'));
 
     await waitFor(() => {
-      expect(enableDatafeedMock).toHaveBeenCalledWith(job, job.latestTimestampMs, true);
+      expect(enableDatafeedMock).toHaveBeenCalledWith(job, job.latestTimestampMs);
       expect(onJobEnabledMock).toHaveBeenCalledWith(job);
     });
   });

diff --git a/..._solution/public/overview/components/entity_analytics/anomalies/components/enable_job.tsx b/..._solution/public/overview/components/entity_analytics/anomalies/components/enable_job.tsx
@@ -24,7 +24,7 @@ export const EnableJob = ({
   const { enableDatafeed, isLoading: isEnabling } = useEnableDataFeed();
 
   const handleChange = useCallback(async () => {
-    const result = await enableDatafeed(job, job.latestTimestampMs || 0, true);
+    const result = await enableDatafeed(job, job.latestTimestampMs || 0);
 
     if (result.enabled) {
       onJobEnabled(job);