[ML] Fix PDF and PNG export with ML embeddables (#128897)

* set up renderComplete callbacks from the swim lane embeddable

* set up renderComplete callbacks from the anomaly charts embeddable

* update output

* set attribute

* update jest tests

x-pack/plugins/ml/public/embeddables/anomaly_charts/anomaly_charts_embeddable.tsx

@@ -94,10 +94,28 @@ export class AnomalyChartsEmbeddable extends Embeddable<
     }
   }
 
+  public onLoading() {
+    this.renderComplete.dispatchInProgress();
+    this.updateOutput({ loading: true, error: undefined });
+  }
+
+  public onError(error: Error) {
+    this.renderComplete.dispatchError();
+    this.updateOutput({ loading: false, error: { name: error.name, message: error.message } });
+  }
+
+  public onRenderComplete() {
+    this.renderComplete.dispatchComplete();
+    this.updateOutput({ loading: false, error: undefined });
+  }
+
   public render(node: HTMLElement) {
     super.render(node);
     this.node = node;
 
+    // required for the export feature to work
+    this.node.setAttribute('data-shared-item', '');
+
     const I18nContext = this.services[0].i18n.Context;
     const theme$ = this.services[0].theme.theme$;
 
@@ -114,6 +132,9 @@ export class AnomalyChartsEmbeddable extends Embeddable<
                 refresh={this.reload$.asObservable()}
                 onInputChange={this.updateInput.bind(this)}
                 onOutputChange={this.updateOutput.bind(this)}
+                onRenderComplete={this.onRenderComplete.bind(this)}
+                onLoading={this.onLoading.bind(this)}
+                onError={this.onError.bind(this)}
               />
             </Suspense>
           </KibanaContextProvider>

x-pack/plugins/ml/public/embeddables/anomaly_charts/embeddable_anomaly_charts_container.test.tsx

@@ -49,6 +49,9 @@ describe('EmbeddableAnomalyChartsContainer', () => {
 
   const onInputChange = jest.fn();
   const onOutputChange = jest.fn();
+  const onRenderComplete = jest.fn();
+  const onLoading = jest.fn();
+  const onError = jest.fn();
 
   const mockedInput = {
     viewMode: 'view',
@@ -145,6 +148,9 @@ describe('EmbeddableAnomalyChartsContainer', () => {
         refresh={refresh}
         onInputChange={onInputChange}
         onOutputChange={onOutputChange}
+        onLoading={onLoading}
+        onRenderComplete={onRenderComplete}
+        onError={onError}
       />,
       defaultOptions
     );
@@ -172,6 +178,9 @@ describe('EmbeddableAnomalyChartsContainer', () => {
         refresh={refresh}
         onInputChange={onInputChange}
         onOutputChange={onOutputChange}
+        onLoading={onLoading}
+        onRenderComplete={onRenderComplete}
+        onError={onError}
       />,
       defaultOptions
     );

x-pack/plugins/ml/public/embeddables/anomaly_charts/embeddable_anomaly_charts_container.tsx

@@ -38,6 +38,9 @@ export interface EmbeddableAnomalyChartsContainerProps {
   refresh: Observable<any>;
   onInputChange: (input: Partial<AnomalyChartsEmbeddableInput>) => void;
   onOutputChange: (output: Partial<AnomalyChartsEmbeddableOutput>) => void;
+  onRenderComplete: () => void;
+  onLoading: () => void;
+  onError: (error: Error) => void;
 }
 
 export const EmbeddableAnomalyChartsContainer: FC<EmbeddableAnomalyChartsContainerProps> = ({
@@ -48,6 +51,9 @@ export const EmbeddableAnomalyChartsContainer: FC<EmbeddableAnomalyChartsContain
   refresh,
   onInputChange,
   onOutputChange,
+  onRenderComplete,
+  onError,
+  onLoading,
 }) => {
   const [chartWidth, setChartWidth] = useState<number>(0);
   const [severity, setSeverity] = useState(
@@ -94,7 +100,8 @@ export const EmbeddableAnomalyChartsContainer: FC<EmbeddableAnomalyChartsContain
     refresh,
     services,
     chartWidth,
-    severity.val
+    severity.val,
+    { onRenderComplete, onError, onLoading }
   );
   const resizeHandler = useCallback(
     throttle((e: { width: number; height: number }) => {

x-pack/plugins/ml/public/embeddables/anomaly_charts/use_anomaly_charts_input_resolver.test.ts

@@ -40,6 +40,12 @@ describe('useAnomalyChartsInputResolver', () => {
   const start = moment().subtract(1, 'years');
   const end = moment();
 
+  const renderCallbacks = {
+    onRenderComplete: jest.fn(),
+    onLoading: jest.fn(),
+    onError: jest.fn(),
+  };
+
   beforeEach(() => {
     jest.useFakeTimers();
 
@@ -116,21 +122,27 @@ describe('useAnomalyChartsInputResolver', () => {
         refresh,
         services,
         1000,
-        0
+        0,
+        renderCallbacks
       )
     );
 
     expect(result.current.chartsData).toBe(undefined);
     expect(result.current.error).toBe(undefined);
     expect(result.current.isLoading).toBe(true);
+    expect(renderCallbacks.onLoading).toHaveBeenCalledTimes(0);
 
     jest.advanceTimersByTime(501);
 
+    expect(renderCallbacks.onLoading).toHaveBeenCalledTimes(1);
+
     const explorerServices = services[2];
 
     expect(explorerServices.anomalyDetectorService.getJobs$).toHaveBeenCalledTimes(1);
     expect(explorerServices.anomalyExplorerService.getAnomalyData$).toHaveBeenCalledTimes(1);
 
+    expect(renderCallbacks.onRenderComplete).toHaveBeenCalledTimes(1);
+
     embeddableInput.next({
       id: 'test-explorer-charts-embeddable',
       jobIds: ['anotherJobId'],
@@ -144,8 +156,14 @@ describe('useAnomalyChartsInputResolver', () => {
     });
     jest.advanceTimersByTime(501);
 
+    expect(renderCallbacks.onLoading).toHaveBeenCalledTimes(2);
+
     expect(explorerServices.anomalyDetectorService.getJobs$).toHaveBeenCalledTimes(2);
     expect(explorerServices.anomalyExplorerService.getAnomalyData$).toHaveBeenCalledTimes(2);
+
+    expect(renderCallbacks.onRenderComplete).toHaveBeenCalledTimes(2);
+
+    expect(renderCallbacks.onError).toHaveBeenCalledTimes(0);
   });
 
   test.skip('should not complete the observable on error', async () => {
@@ -156,7 +174,8 @@ describe('useAnomalyChartsInputResolver', () => {
         refresh,
         services,
         1000,
-        1
+        1,
+        renderCallbacks
       )
     );
 
@@ -168,5 +187,6 @@ describe('useAnomalyChartsInputResolver', () => {
     } as Partial<AnomalyChartsEmbeddableInput>);
 
     expect(result.current.error).toBeDefined();
+    expect(renderCallbacks.onError).toHaveBeenCalledTimes(1);
   });
 });

x-pack/plugins/ml/public/embeddables/anomaly_charts/use_anomaly_charts_input_resolver.ts

@@ -35,7 +35,12 @@ export function useAnomalyChartsInputResolver(
   refresh: Observable<any>,
   services: [CoreStart, MlStartDependencies, AnomalyChartsServices],
   chartWidth: number,
-  severity: number
+  severity: number,
+  renderCallbacks: {
+    onRenderComplete: () => void;
+    onLoading: () => void;
+    onError: (error: Error) => void;
+  }
 ): {
   chartsData: ExplorerChartsData | undefined;
   isLoading: boolean;
@@ -61,6 +66,9 @@ export function useAnomalyChartsInputResolver(
       .pipe(
         tap(setIsLoading.bind(null, true)),
         debounceTime(FETCH_RESULTS_DEBOUNCE_MS),
+        tap(() => {
+          renderCallbacks.onLoading();
+        }),
         switchMap(([explorerJobs, input, embeddableContainerWidth, severityValue]) => {
           if (!explorerJobs) {
             // couldn't load the list of jobs
@@ -118,6 +126,8 @@ export function useAnomalyChartsInputResolver(
           setError(null);
           setChartsData(results);
           setIsLoading(false);
+
+          renderCallbacks.onRenderComplete();
         }
       });
 
@@ -134,5 +144,11 @@ export function useAnomalyChartsInputResolver(
     severity$.next(severity);
   }, [severity]);
 
+  useEffect(() => {
+    if (error) {
+      renderCallbacks.onError(error);
+    }
+  }, [error]);
+
   return { chartsData, isLoading, error };
 }

x-pack/plugins/ml/public/embeddables/anomaly_swimlane/anomaly_swimlane_embeddable.tsx

@@ -56,10 +56,28 @@ export class AnomalySwimlaneEmbeddable extends Embeddable<
     );
   }
 
+  public onLoading() {
+    this.renderComplete.dispatchInProgress();
+    this.updateOutput({ loading: true, error: undefined });
+  }
+
+  public onError(error: Error) {
+    this.renderComplete.dispatchError();
+    this.updateOutput({ loading: false, error: { name: error.name, message: error.message } });
+  }
+
+  public onRenderComplete() {
+    this.renderComplete.dispatchComplete();
+    this.updateOutput({ loading: false, error: undefined });
+  }
+
   public render(node: HTMLElement) {
     super.render(node);
     this.node = node;
 
+    // required for the export feature to work
+    this.node.setAttribute('data-shared-item', '');
+
     const I18nContext = this.services[0].i18n.Context;
     const theme$ = this.services[0].theme.theme$;
 
@@ -76,6 +94,9 @@ export class AnomalySwimlaneEmbeddable extends Embeddable<
                 refresh={this.reload$.asObservable()}
                 onInputChange={this.updateInput.bind(this)}
                 onOutputChange={this.updateOutput.bind(this)}
+                onRenderComplete={this.onRenderComplete.bind(this)}
+                onLoading={this.onLoading.bind(this)}
+                onError={this.onError.bind(this)}
               />
             </Suspense>
           </KibanaContextProvider>

x-pack/plugins/ml/public/embeddables/anomaly_swimlane/embeddable_swim_lane_container.test.tsx

@@ -48,6 +48,9 @@ describe('ExplorerSwimlaneContainer', () => {
 
   const onInputChange = jest.fn();
   const onOutputChange = jest.fn();
+  const onRenderComplete = jest.fn();
+  const onLoading = jest.fn();
+  const onError = jest.fn();
 
   beforeEach(() => {
     embeddableContext = { id: 'test-id' } as AnomalySwimlaneEmbeddable;
@@ -102,6 +105,9 @@ describe('ExplorerSwimlaneContainer', () => {
         refresh={refresh}
         onInputChange={onInputChange}
         onOutputChange={onOutputChange}
+        onLoading={onLoading}
+        onRenderComplete={onRenderComplete}
+        onError={onError}
       />,
       defaultOptions
     );
@@ -141,6 +147,9 @@ describe('ExplorerSwimlaneContainer', () => {
         refresh={refresh}
         onInputChange={onInputChange}
         onOutputChange={onOutputChange}
+        onLoading={onLoading}
+        onRenderComplete={onRenderComplete}
+        onError={onError}
       />,
       defaultOptions
     );

x-pack/plugins/ml/public/embeddables/anomaly_swimlane/embeddable_swim_lane_container.tsx

@@ -35,6 +35,9 @@ export interface ExplorerSwimlaneContainerProps {
   refresh: Observable<any>;
   onInputChange: (input: Partial<AnomalySwimlaneEmbeddableInput>) => void;
   onOutputChange: (output: Partial<AnomalySwimlaneEmbeddableOutput>) => void;
+  onRenderComplete: () => void;
+  onLoading: () => void;
+  onError: (error: Error) => void;
 }
 
 export const EmbeddableSwimLaneContainer: FC<ExplorerSwimlaneContainerProps> = ({
@@ -45,6 +48,9 @@ export const EmbeddableSwimLaneContainer: FC<ExplorerSwimlaneContainerProps> = (
   refresh,
   onInputChange,
   onOutputChange,
+  onRenderComplete,
+  onLoading,
+  onError,
 }) => {
   const [chartWidth, setChartWidth] = useState<number>(0);
 
@@ -61,7 +67,8 @@ export const EmbeddableSwimLaneContainer: FC<ExplorerSwimlaneContainerProps> = (
       refresh,
       services,
       chartWidth,
-      fromPage
+      fromPage,
+      { onRenderComplete, onError, onLoading }
     );
 
   useEffect(() => {

x-pack/plugins/ml/public/embeddables/anomaly_swimlane/swimlane_input_resolver.test.ts

@@ -19,6 +19,12 @@ describe('useSwimlaneInputResolver', () => {
   let services: [CoreStart, MlStartDependencies, AnomalySwimlaneServices];
   let onInputChange: jest.Mock;
 
+  const renderCallbacks = {
+    onRenderComplete: jest.fn(),
+    onLoading: jest.fn(),
+    onError: jest.fn(),
+  };
+
   beforeEach(() => {
     jest.useFakeTimers();
 
@@ -78,6 +84,7 @@ describe('useSwimlaneInputResolver', () => {
     ];
     onInputChange = jest.fn();
   });
+
   afterEach(() => {
     jest.useRealTimers();
     jest.clearAllMocks();
@@ -91,7 +98,8 @@ describe('useSwimlaneInputResolver', () => {
         refresh,
         services,
         1000,
-        1
+        1,
+        renderCallbacks
       )
     );
 
@@ -106,6 +114,9 @@ describe('useSwimlaneInputResolver', () => {
     expect(services[2].anomalyDetectorService.getJobs$).toHaveBeenCalledTimes(1);
     expect(services[2].anomalyTimelineService.loadOverallData).toHaveBeenCalledTimes(1);
 
+    expect(renderCallbacks.onLoading).toHaveBeenCalledTimes(1);
+    expect(renderCallbacks.onRenderComplete).toHaveBeenCalledTimes(1);
+
     await act(async () => {
       embeddableInput.next({
         id: 'test-swimlane-embeddable',
@@ -121,6 +132,9 @@ describe('useSwimlaneInputResolver', () => {
     expect(services[2].anomalyDetectorService.getJobs$).toHaveBeenCalledTimes(2);
     expect(services[2].anomalyTimelineService.loadOverallData).toHaveBeenCalledTimes(2);
 
+    expect(renderCallbacks.onLoading).toHaveBeenCalledTimes(2);
+    expect(renderCallbacks.onRenderComplete).toHaveBeenCalledTimes(2);
+
     await act(async () => {
       embeddableInput.next({
         id: 'test-swimlane-embeddable',
@@ -135,6 +149,9 @@ describe('useSwimlaneInputResolver', () => {
 
     expect(services[2].anomalyDetectorService.getJobs$).toHaveBeenCalledTimes(2);
     expect(services[2].anomalyTimelineService.loadOverallData).toHaveBeenCalledTimes(3);
+
+    expect(renderCallbacks.onLoading).toHaveBeenCalledTimes(3);
+    expect(renderCallbacks.onRenderComplete).toHaveBeenCalledTimes(3);
   });
 
   test('should not complete the observable on error', async () => {
@@ -145,7 +162,8 @@ describe('useSwimlaneInputResolver', () => {
         refresh,
         services,
         1000,
-        1
+        1,
+        renderCallbacks
       )
     );
 
@@ -160,5 +178,7 @@ describe('useSwimlaneInputResolver', () => {
     });
 
     expect(result.current[6]?.message).toBe('Invalid job');
+
+    expect(renderCallbacks.onError).toHaveBeenCalledTimes(1);
   });
 });

x-pack/plugins/ml/public/embeddables/anomaly_swimlane/swimlane_input_resolver.ts

@@ -46,10 +46,15 @@ const FETCH_RESULTS_DEBOUNCE_MS = 500;
 export function useSwimlaneInputResolver(
   embeddableInput$: Observable<AnomalySwimlaneEmbeddableInput>,
   onInputChange: (output: Partial<AnomalySwimlaneEmbeddableOutput>) => void,
-  refresh: Observable<any>,
+  refresh: Observable<void>,
   services: [CoreStart, MlStartDependencies, AnomalySwimlaneServices],
   chartWidth: number,
-  fromPage: number
+  fromPage: number,
+  renderCallbacks: {
+    onRenderComplete: () => void;
+    onLoading: () => void;
+    onError: (error: Error) => void;
+  }
 ): [
   string | undefined,
   OverallSwimlaneData | undefined,
@@ -122,6 +127,9 @@ export function useSwimlaneInputResolver(
       .pipe(
         tap(setIsLoading.bind(null, true)),
         debounceTime(FETCH_RESULTS_DEBOUNCE_MS),
+        tap(() => {
+          renderCallbacks.onLoading();
+        }),
         switchMap(([explorerJobs, input, bucketInterval, fromPageInput, perPageFromState]) => {
           if (!explorerJobs) {
             // couldn't load the list of jobs
@@ -227,6 +235,18 @@ export function useSwimlaneInputResolver(
     chartWidth$.next(chartWidth);
   }, [chartWidth]);
 
+  useEffect(() => {
+    if (error) {
+      renderCallbacks.onError(error);
+    }
+  }, [error]);
+
+  useEffect(() => {
+    if (swimlaneData) {
+      renderCallbacks.onRenderComplete();
+    }
+  }, [swimlaneData]);
+
   return [
     swimlaneType,
     swimlaneData,