[Lens] Use suggestion system in chart switcher for subtypes (#64613)

Co-authored-by: Elastic Machine <elasticmachine@users.noreply.github.com>

Author: Wylie Conlon

x-pack/plugins/lens/public/datatable_visualization/visualization.tsx

@@ -41,6 +41,10 @@ export const datatableVisualization: Visualization<
     },
   ],
 
+  getVisualizationTypeId() {
+    return 'lnsDatatable';
+  },
+
   getLayerIds(state) {
     return state.layers.map(l => l.layerId);
   },

x-pack/plugins/lens/public/editor_frame_service/editor_frame/config_panel/chart_switch.test.tsx

@@ -62,7 +62,25 @@ describe('chart_switch', () => {
             id: 'subvisC2',
             label: 'C2',
           },
+          {
+            icon: 'empty',
+            id: 'subvisC3',
+            label: 'C3',
+          },
         ],
+        getSuggestions: jest.fn(options => {
+          if (options.subVisualizationId === 'subvisC2') {
+            return [];
+          }
+          return [
+            {
+              score: 1,
+              title: '',
+              state: `suggestion`,
+              previewIcon: 'empty',
+            },
+          ];
+        }),
       },
     };
   }
@@ -313,10 +331,11 @@ describe('chart_switch', () => {
     expect(getMenuItem('subvisB', component).prop('betaBadgeIconType')).toBeUndefined();
   });
 
-  it('should not indicate data loss if visualization is not changed', () => {
+  it('should not show a warning when the subvisualization is the same', () => {
     const dispatch = jest.fn();
     const frame = mockFrame(['a', 'b', 'c']);
     const visualizations = mockVisualizations();
+    visualizations.visC.getVisualizationTypeId.mockReturnValue('subvisC2');
     const switchVisualizationType = jest.fn(() => 'therebedragons');
 
     visualizations.visC.switchVisualizationType = switchVisualizationType;
@@ -333,10 +352,10 @@ describe('chart_switch', () => {
       />
     );
 
-    expect(getMenuItem('subvisC2', component).prop('betaBadgeIconType')).toBeUndefined();
+    expect(getMenuItem('subvisC2', component).prop('betaBadgeIconType')).not.toBeDefined();
   });
 
-  it('should remove all layers if there is no suggestion', () => {
+  it('should get suggestions when switching subvisualization', () => {
     const dispatch = jest.fn();
     const visualizations = mockVisualizations();
     visualizations.visB.getSuggestions.mockReturnValueOnce([]);
@@ -377,7 +396,7 @@ describe('chart_switch', () => {
     const dispatch = jest.fn();
     const frame = mockFrame(['a', 'b', 'c']);
     const visualizations = mockVisualizations();
-    const switchVisualizationType = jest.fn(() => 'therebedragons');
+    const switchVisualizationType = jest.fn(() => 'switched');
 
     visualizations.visC.switchVisualizationType = switchVisualizationType;
 
@@ -393,12 +412,12 @@ describe('chart_switch', () => {
       />
     );
 
-    switchTo('subvisC2', component);
-    expect(switchVisualizationType).toHaveBeenCalledWith('subvisC2', 'therebegriffins');
+    switchTo('subvisC3', component);
+    expect(switchVisualizationType).toHaveBeenCalledWith('subvisC3', 'suggestion');
     expect(dispatch).toHaveBeenCalledWith(
       expect.objectContaining({
         type: 'SWITCH_VISUALIZATION',
-        initialState: 'therebedragons',
+        initialState: 'switched',
       })
     );
     expect(frame.removeLayers).not.toHaveBeenCalled();

x-pack/plugins/lens/public/editor_frame_service/editor_frame/config_panel/chart_switch.tsx

@@ -105,7 +105,16 @@ export function ChartSwitch(props: Props) {
     const switchVisType =
       props.visualizationMap[visualizationId].switchVisualizationType ||
       ((_type: string, initialState: unknown) => initialState);
-    if (props.visualizationId === visualizationId) {
+    const layers = Object.entries(props.framePublicAPI.datasourceLayers);
+    const containsData = layers.some(
+      ([_layerId, datasource]) => datasource.getTableSpec().length > 0
+    );
+    // Always show the active visualization as a valid selection
+    if (
+      props.visualizationId === visualizationId &&
+      props.visualizationState &&
+      newVisualization.getVisualizationTypeId(props.visualizationState) === subVisualizationId
+    ) {
       return {
         visualizationId,
         subVisualizationId,
@@ -116,13 +125,13 @@ export function ChartSwitch(props: Props) {
       };
     }
 
-    const layers = Object.entries(props.framePublicAPI.datasourceLayers);
-    const containsData = layers.some(
-      ([_layerId, datasource]) => datasource.getTableSpec().length > 0
+    const topSuggestion = getTopSuggestion(
+      props,
+      visualizationId,
+      newVisualization,
+      subVisualizationId
     );
 
-    const topSuggestion = getTopSuggestion(props, visualizationId, newVisualization);
-
     let dataLoss: VisualizationSelection['dataLoss'];
 
     if (!containsData) {
@@ -250,14 +259,16 @@ export function ChartSwitch(props: Props) {
 function getTopSuggestion(
   props: Props,
   visualizationId: string,
-  newVisualization: Visualization<unknown, unknown>
+  newVisualization: Visualization<unknown, unknown>,
+  subVisualizationId?: string
 ): Suggestion | undefined {
   const suggestions = getSuggestions({
     datasourceMap: props.datasourceMap,
     datasourceStates: props.datasourceStates,
     visualizationMap: { [visualizationId]: newVisualization },
     activeVisualizationId: props.visualizationId,
     visualizationState: props.visualizationState,
+    subVisualizationId,
   }).filter(suggestion => {
     // don't use extended versions of current data table on switching between visualizations
     // to avoid confusing the user.

x-pack/plugins/lens/public/editor_frame_service/editor_frame/suggestion_helpers.ts

@@ -44,6 +44,7 @@ export function getSuggestions({
   datasourceStates,
   visualizationMap,
   activeVisualizationId,
+  subVisualizationId,
   visualizationState,
   field,
 }: {
@@ -57,6 +58,7 @@ export function getSuggestions({
   >;
   visualizationMap: Record<string, Visualization>;
   activeVisualizationId: string | null;
+  subVisualizationId?: string;
   visualizationState: unknown;
   field?: unknown;
 }): Suggestion[] {
@@ -89,7 +91,8 @@ export function getSuggestions({
             table,
             visualizationId,
             datasourceSuggestion,
-            currentVisualizationState
+            currentVisualizationState,
+            subVisualizationId
           );
         })
       )
@@ -108,13 +111,15 @@ function getVisualizationSuggestions(
   table: TableSuggestion,
   visualizationId: string,
   datasourceSuggestion: DatasourceSuggestion & { datasourceId: string },
-  currentVisualizationState: unknown
+  currentVisualizationState: unknown,
+  subVisualizationId?: string
 ) {
   return visualization
     .getSuggestions({
       table,
       state: currentVisualizationState,
       keptLayerIds: datasourceSuggestion.keptLayerIds,
+      subVisualizationId,
     })
     .map(({ state, ...visualizationSuggestion }) => ({
       ...visualizationSuggestion,

x-pack/plugins/lens/public/editor_frame_service/mocks.tsx

@@ -28,6 +28,7 @@ export function createMockVisualization(): jest.Mocked<Visualization> {
         label: 'TEST',
       },
     ],
+    getVisualizationTypeId: jest.fn(_state => 'empty'),
     getDescription: jest.fn(_state => ({ label: '' })),
     switchVisualizationType: jest.fn((_, x) => x),
     getPersistableState: jest.fn(_state => _state),

x-pack/plugins/lens/public/metric_visualization/metric_visualization.tsx

@@ -53,6 +53,10 @@ export const metricVisualization: Visualization<State, PersistableState> = {
     },
   ],
 
+  getVisualizationTypeId() {
+    return 'lnsMetric';
+  },
+
   clearLayer(state) {
     return {
       ...state,

x-pack/plugins/lens/public/types.ts

@@ -312,6 +312,10 @@ export interface SuggestionRequest<T = unknown> {
    * The visualization needs to know which table is being suggested
    */
   keptLayerIds: string[];
+  /**
+   * Different suggestions can be generated for each subtype of the visualization
+   */
+  subVisualizationId?: string;
 }
 
 /**
@@ -388,6 +392,11 @@ export interface Visualization<T = unknown, P = unknown> {
    * but can register multiple subtypes
    */
   visualizationTypes: VisualizationType[];
+  /**
+   * Return the ID of the current visualization. Used to highlight
+   * the active subtype of the visualization.
+   */
+  getVisualizationTypeId: (state: T) => string;
   /**
    * If the visualization has subtypes, update the subtype in state.
    */

x-pack/plugins/lens/public/xy_visualization/xy_visualization.test.ts

@@ -27,7 +27,7 @@ function exampleState(): State {
 }
 
 describe('xy_visualization', () => {
-  describe('getDescription', () => {
+  describe('#getDescription', () => {
     function mixedState(...types: SeriesType[]) {
       const state = exampleState();
       return {
@@ -81,6 +81,45 @@ describe('xy_visualization', () => {
     });
   });
 
+  describe('#getVisualizationTypeId', () => {
+    function mixedState(...types: SeriesType[]) {
+      const state = exampleState();
+      return {
+        ...state,
+        layers: types.map((t, i) => ({
+          ...state.layers[0],
+          layerId: `layer_${i}`,
+          seriesType: t,
+        })),
+      };
+    }
+
+    it('should show mixed when each layer is different', () => {
+      expect(xyVisualization.getVisualizationTypeId(mixedState('bar', 'line'))).toEqual('mixed');
+    });
+
+    it('should show the preferredSeriesType if there are no layers', () => {
+      expect(xyVisualization.getVisualizationTypeId(mixedState())).toEqual('bar');
+    });
+
+    it('should combine multiple layers into one type', () => {
+      expect(
+        xyVisualization.getVisualizationTypeId(mixedState('bar_horizontal', 'bar_horizontal'))
+      ).toEqual('bar_horizontal');
+    });
+
+    it('should return the subtype for single layers', () => {
+      expect(xyVisualization.getVisualizationTypeId(mixedState('area'))).toEqual('area');
+      expect(xyVisualization.getVisualizationTypeId(mixedState('line'))).toEqual('line');
+      expect(xyVisualization.getVisualizationTypeId(mixedState('area_stacked'))).toEqual(
+        'area_stacked'
+      );
+      expect(xyVisualization.getVisualizationTypeId(mixedState('bar_horizontal_stacked'))).toEqual(
+        'bar_horizontal_stacked'
+      );
+    });
+  });
+
   describe('#initialize', () => {
     it('loads default state', () => {
       const mockFrame = createMockFramePublicAPI();

x-pack/plugins/lens/public/xy_visualization/xy_visualization.tsx

@@ -12,7 +12,7 @@ import { I18nProvider } from '@kbn/i18n/react';
 import { i18n } from '@kbn/i18n';
 import { getSuggestions } from './xy_suggestions';
 import { LayerContextMenu } from './xy_config_panel';
-import { Visualization, OperationMetadata } from '../types';
+import { Visualization, OperationMetadata, VisualizationType } from '../types';
 import { State, PersistableState, SeriesType, visualizationTypes, LayerConfig } from './types';
 import { toExpression, toPreviewExpression } from './to_expression';
 import chartBarStackedSVG from '../assets/chart_bar_stacked.svg';
@@ -24,6 +24,18 @@ const defaultSeriesType = 'bar_stacked';
 const isNumericMetric = (op: OperationMetadata) => !op.isBucketed && op.dataType === 'number';
 const isBucketed = (op: OperationMetadata) => op.isBucketed;
 
+function getVisualizationType(state: State): VisualizationType | 'mixed' {
+  if (!state.layers.length) {
+    return (
+      visualizationTypes.find(t => t.id === state.preferredSeriesType) ?? visualizationTypes[0]
+    );
+  }
+  const visualizationType = visualizationTypes.find(t => t.id === state.layers[0].seriesType);
+  const seriesTypes = _.unique(state.layers.map(l => l.seriesType));
+
+  return visualizationType && seriesTypes.length === 1 ? visualizationType : 'mixed';
+}
+
 function getDescription(state?: State) {
   if (!state) {
     return {
@@ -34,39 +46,42 @@ function getDescription(state?: State) {
     };
   }
 
+  const visualizationType = getVisualizationType(state);
+
   if (!state.layers.length) {
-    const visualizationType = visualizationTypes.find(v => v.id === state.preferredSeriesType)!;
+    const preferredType = visualizationType as VisualizationType;
     return {
-      icon: visualizationType.largeIcon || visualizationType.icon,
-      label: visualizationType.label,
+      icon: preferredType.largeIcon || preferredType.icon,
+      label: preferredType.label,
     };
   }
 
-  const visualizationType = visualizationTypes.find(t => t.id === state.layers[0].seriesType)!;
-  const seriesTypes = _.unique(state.layers.map(l => l.seriesType));
-
   return {
     icon:
-      seriesTypes.length === 1
-        ? visualizationType.largeIcon || visualizationType.icon
-        : chartMixedSVG,
+      visualizationType === 'mixed'
+        ? chartMixedSVG
+        : visualizationType.largeIcon || visualizationType.icon,
     label:
-      seriesTypes.length === 1
-        ? visualizationType.label
-        : isHorizontalChart(state.layers)
-        ? i18n.translate('xpack.lens.xyVisualization.mixedBarHorizontalLabel', {
-            defaultMessage: 'Mixed horizontal bar',
-          })
-        : i18n.translate('xpack.lens.xyVisualization.mixedLabel', {
-            defaultMessage: 'Mixed XY',
-          }),
+      visualizationType === 'mixed'
+        ? isHorizontalChart(state.layers)
+          ? i18n.translate('xpack.lens.xyVisualization.mixedBarHorizontalLabel', {
+              defaultMessage: 'Mixed horizontal bar',
+            })
+          : i18n.translate('xpack.lens.xyVisualization.mixedLabel', {
+              defaultMessage: 'Mixed XY',
+            })
+        : visualizationType.label,
   };
 }
 
 export const xyVisualization: Visualization<State, PersistableState> = {
   id: 'lnsXY',
 
   visualizationTypes,
+  getVisualizationTypeId(state) {
+    const type = getVisualizationType(state);
+    return type === 'mixed' ? type : type.id;
+  },
 
   getLayerIds(state) {
     return state.layers.map(l => l.layerId);