[Lens] Use suggestion system in chart switcher for subtypes

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

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

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

@@ -58,7 +58,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',
+            },
+          ];
+        }),
       },
     };
   }
@@ -309,7 +327,31 @@ 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;
+
+    const component = mount(
+      <ChartSwitch
+        visualizationId="visC"
+        visualizationState={'therebegriffins'}
+        visualizationMap={visualizations}
+        dispatch={dispatch}
+        framePublicAPI={frame}
+        datasourceMap={mockDatasourceMap()}
+        datasourceStates={mockDatasourceStates()}
+      />
+    );
+
+    expect(getMenuItem('subvisC2', component).prop('betaBadgeIconType')).not.toBeDefined();
+  });
+
+  it('should get suggestions when switching subvisualization', () => {
     const dispatch = jest.fn();
     const frame = mockFrame(['a', 'b', 'c']);
     const visualizations = mockVisualizations();
@@ -329,7 +371,7 @@ describe('chart_switch', () => {
       />
     );
 
-    expect(getMenuItem('subvisC2', component).prop('betaBadgeIconType')).toBeUndefined();
+    expect(getMenuItem('subvisC2', component).prop('betaBadgeIconType')).toEqual('alert');
   });
 
   it('should remove all layers if there is no suggestion', () => {
@@ -369,11 +411,11 @@ describe('chart_switch', () => {
     );
   });
 
-  it('should not remove layers when switching between subtypes', () => {
+  it('should not remove layers when switching between subtypes with valid suggestions', () => {
     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;
 
@@ -389,12 +431,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/legacy/plugins/lens/public/editor_frame_service/editor_frame/chart_switch.tsx

@@ -105,7 +105,17 @@ 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 +126,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 +260,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/legacy/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/legacy/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/legacy/plugins/lens/public/metric_visualization/metric_visualization.tsx

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

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

@@ -315,6 +315,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;
 }
 
 /**
@@ -371,9 +375,21 @@ export interface VisualizationType {
 }
 
 export interface Visualization<T = unknown, P = unknown> {
+  /**
+   * Visualization ID for internal Lens use. Not shown to users.
+   */
   id: string;
 
+  /**
+   * Visualizations can provide multiple subtypes with unique IDs
+   * that can share the same state, for example bar/line/area.
+   */
   visualizationTypes: VisualizationType[];
+  /**
+   * Return the ID of the current visualization. Used to highlight
+   * the active subtype of the visualization.
+   */
+  getVisualizationTypeId: (state: T) => string;
 
   getLayerIds: (state: T) => string[];
   clearLayer: (state: T, layerId: string) => T;

x-pack/legacy/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();