[Lens] Fix switching with layers (elastic#71982)

* [Lens] Fix chart switching with multiple layers

* Unskip Lens smokescreen test

* Fix types

* Revert <p> change

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

@@ -129,7 +129,7 @@ function LayerPanels(
               className="lnsConfigPanel__addLayerBtn"
               fullWidth
               size="s"
-              data-test-subj="lnsXY_layer_add"
+              data-test-subj="lnsLayerAddButton"
               aria-label={i18n.translate('xpack.lens.xyChart.addLayerButton', {
                 defaultMessage: 'Add layer',
               })}

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

@@ -93,14 +93,14 @@ describe('LayerPanel', () => {
   describe('layer reset and remove', () => {
     it('should show the reset button when single layer', () => {
       const component = mountWithIntl(<LayerPanel {...getDefaultProps()} />);
-      expect(component.find('[data-test-subj="lns_layer_remove"]').first().text()).toContain(
+      expect(component.find('[data-test-subj="lnsLayerRemove"]').first().text()).toContain(
         'Reset layer'
       );
     });
 
     it('should show the delete button when multiple layers', () => {
       const component = mountWithIntl(<LayerPanel {...getDefaultProps()} isOnlyLayer={false} />);
-      expect(component.find('[data-test-subj="lns_layer_remove"]').first().text()).toContain(
+      expect(component.find('[data-test-subj="lnsLayerRemove"]').first().text()).toContain(
         'Delete layer'
       );
     });
@@ -109,7 +109,7 @@ describe('LayerPanel', () => {
       const cb = jest.fn();
       const component = mountWithIntl(<LayerPanel {...getDefaultProps()} onRemoveLayer={cb} />);
       act(() => {
-        component.find('[data-test-subj="lns_layer_remove"]').first().simulate('click');
+        component.find('[data-test-subj="lnsLayerRemove"]').first().simulate('click');
       });
       expect(cb).toHaveBeenCalled();
     });

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

@@ -429,7 +429,7 @@ export function LayerPanel(
               size="xs"
               iconType="trash"
               color="danger"
-              data-test-subj="lns_layer_remove"
+              data-test-subj="lnsLayerRemove"
               onClick={() => {
                 // If we don't blur the remove / clear button, it remains focused
                 // which is a strange UX in this case. e.target.blur doesn't work

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

@@ -46,6 +46,16 @@ describe('chart_switch', () => {
     };
   }
 
+  /**
+   * There are three visualizations. Each one has the same suggestion behavior:
+   *
+   * visA: suggests an empty state
+   * visB: suggests an empty state
+   * visC:
+   *  - Never switches to subvisC2
+   *  - Allows a switch to subvisC3
+   *  - Allows a switch to subvisC1
+   */
   function mockVisualizations() {
     return {
       visA: generateVisualization('visA'),
@@ -292,6 +302,49 @@ describe('chart_switch', () => {
     expect(getMenuItem('visB', component).prop('betaBadgeIconType')).toEqual('alert');
   });
 
+  it('should support multi-layer suggestions without data loss', () => {
+    const dispatch = jest.fn();
+    const visualizations = mockVisualizations();
+    const frame = mockFrame(['a', 'b']);
+
+    const datasourceMap = mockDatasourceMap();
+    datasourceMap.testDatasource.getDatasourceSuggestionsFromCurrentState.mockReturnValue([
+      {
+        state: {},
+        table: {
+          columns: [
+            {
+              columnId: 'a',
+              operation: {
+                label: '',
+                dataType: 'string',
+                isBucketed: true,
+              },
+            },
+          ],
+          isMultiRow: true,
+          layerId: 'a',
+          changeType: 'unchanged',
+        },
+        keptLayerIds: ['a', 'b'],
+      },
+    ]);
+
+    const component = mount(
+      <ChartSwitch
+        visualizationId="visA"
+        visualizationState={{}}
+        visualizationMap={visualizations}
+        dispatch={dispatch}
+        framePublicAPI={frame}
+        datasourceMap={datasourceMap}
+        datasourceStates={mockDatasourceStates()}
+      />
+    );
+
+    expect(getMenuItem('visB', component).prop('betaBadgeIconType')).toBeUndefined();
+  });
+
   it('should indicate data loss if no data will be used', () => {
     const dispatch = jest.fn();
     const visualizations = mockVisualizations();

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

@@ -139,7 +139,7 @@ export function ChartSwitch(props: Props) {
       dataLoss = 'nothing';
     } else if (!topSuggestion) {
       dataLoss = 'everything';
-    } else if (layers.length > 1) {
+    } else if (layers.length > 1 && layers.length !== topSuggestion.keptLayerIds.length) {
       dataLoss = 'layers';
     } else if (topSuggestion.columns !== layers[0][1].getTableSpec().length) {
       dataLoss = 'columns';
@@ -258,14 +258,15 @@ function getTopSuggestion(
   newVisualization: Visualization<unknown, unknown>,
   subVisualizationId?: string
 ): Suggestion | undefined {
-  const suggestions = getSuggestions({
+  const unfilteredSuggestions = getSuggestions({
     datasourceMap: props.datasourceMap,
     datasourceStates: props.datasourceStates,
     visualizationMap: { [visualizationId]: newVisualization },
     activeVisualizationId: props.visualizationId,
     visualizationState: props.visualizationState,
     subVisualizationId,
-  }).filter((suggestion) => {
+  });
+  const suggestions = unfilteredSuggestions.filter((suggestion) => {
     // don't use extended versions of current data table on switching between visualizations
     // to avoid confusing the user.
     return (

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

@@ -13,6 +13,7 @@ import {
 } from '../types';
 import { State, XYState, visualizationTypes } from './types';
 import { generateId } from '../id_generator';
+import { xyVisualization } from './xy_visualization';
 
 jest.mock('../id_generator');
 
@@ -119,7 +120,33 @@ describe('xy_suggestions', () => {
     });
 
     expect(suggestions).toHaveLength(visualizationTypes.length);
-    expect(suggestions.map(({ state }) => state.preferredSeriesType)).toEqual([
+    expect(suggestions.map(({ state }) => xyVisualization.getVisualizationTypeId(state))).toEqual([
+      'bar_stacked',
+      'area_stacked',
+      'area',
+      'line',
+      'bar_horizontal_stacked',
+      'bar_horizontal',
+      'bar',
+    ]);
+  });
+
+  // This limitation is acceptable for now, but is now tested
+  test('is unable to generate layers when switching from a non-XY chart with multiple layers', () => {
+    (generateId as jest.Mock).mockReturnValueOnce('aaa');
+    const suggestions = getSuggestions({
+      table: {
+        isMultiRow: true,
+        columns: [numCol('bytes'), dateCol('date')],
+        layerId: 'first',
+        changeType: 'unchanged',
+      },
+      keptLayerIds: ['first', 'second'],
+    });
+
+    expect(suggestions).toHaveLength(visualizationTypes.length);
+    expect(suggestions.map(({ state }) => state.layers.length)).toEqual([1, 1, 1, 1, 1, 1, 1]);
+    expect(suggestions.map(({ state }) => xyVisualization.getVisualizationTypeId(state))).toEqual([
       'bar_stacked',
       'area_stacked',
       'area',
@@ -156,7 +183,51 @@ describe('xy_suggestions', () => {
     });
 
     expect(suggestions).toHaveLength(visualizationTypes.length);
-    expect(suggestions.map(({ state }) => state.preferredSeriesType)).toEqual([
+    expect(suggestions.map(({ state }) => xyVisualization.getVisualizationTypeId(state))).toEqual([
+      'line',
+      'bar',
+      'bar_horizontal',
+      'bar_stacked',
+      'bar_horizontal_stacked',
+      'area',
+      'area_stacked',
+    ]);
+  });
+
+  test('suggests all basic x y charts when switching from another x y chart with multiple layers', () => {
+    (generateId as jest.Mock).mockReturnValueOnce('aaa');
+    const suggestions = getSuggestions({
+      table: {
+        isMultiRow: true,
+        columns: [numCol('bytes'), dateCol('date')],
+        layerId: 'first',
+        changeType: 'unchanged',
+      },
+      keptLayerIds: ['first', 'second'],
+      state: {
+        legend: { isVisible: true, position: 'bottom' },
+        preferredSeriesType: 'bar',
+        layers: [
+          {
+            layerId: 'first',
+            seriesType: 'bar',
+            xAccessor: 'date',
+            accessors: ['bytes'],
+            splitAccessor: undefined,
+          },
+          {
+            layerId: 'second',
+            seriesType: 'bar',
+            xAccessor: undefined,
+            accessors: [],
+            splitAccessor: undefined,
+          },
+        ],
+      },
+    });
+
+    expect(suggestions).toHaveLength(visualizationTypes.length);
+    expect(suggestions.map(({ state }) => xyVisualization.getVisualizationTypeId(state))).toEqual([
       'line',
       'bar',
       'bar_horizontal',
@@ -165,6 +236,15 @@ describe('xy_suggestions', () => {
       'area',
       'area_stacked',
     ]);
+    expect(suggestions.map(({ state }) => state.layers.map((l) => l.layerId))).toEqual([
+      ['first', 'second'],
+      ['first', 'second'],
+      ['first', 'second'],
+      ['first', 'second'],
+      ['first', 'second'],
+      ['first', 'second'],
+      ['first', 'second'],
+    ]);
   });
 
   test('suggests all basic x y chart with date on x', () => {
@@ -388,7 +468,7 @@ describe('xy_suggestions', () => {
         changeType: 'unchanged',
       },
       state: currentState,
-      keptLayerIds: [],
+      keptLayerIds: ['first'],
     });
 
     expect(rest).toHaveLength(visualizationTypes.length - 2);
@@ -497,7 +577,7 @@ describe('xy_suggestions', () => {
         changeType: 'extended',
       },
       state: currentState,
-      keptLayerIds: [],
+      keptLayerIds: ['first'],
     });
 
     expect(rest).toHaveLength(0);
@@ -536,7 +616,7 @@ describe('xy_suggestions', () => {
         changeType: 'reorder',
       },
       state: currentState,
-      keptLayerIds: [],
+      keptLayerIds: ['first'],
     });
 
     expect(rest).toHaveLength(0);
@@ -576,7 +656,7 @@ describe('xy_suggestions', () => {
         changeType: 'extended',
       },
       state: currentState,
-      keptLayerIds: [],
+      keptLayerIds: ['first'],
     });
 
     expect(rest).toHaveLength(0);

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

@@ -394,17 +394,25 @@ function buildSuggestion({
         : undefined,
   };
 
+  // Maintain consistent order for any layers that were saved
   const keptLayers = currentState
-    ? currentState.layers.filter(
-        (layer) => layer.layerId !== layerId && keptLayerIds.includes(layer.layerId)
-      )
+    ? currentState.layers
+        // Remove layers that aren't being suggested
+        .filter((layer) => keptLayerIds.includes(layer.layerId))
+        // Update in place
+        .map((layer) => (layer.layerId === layerId ? newLayer : layer))
+        // Replace the seriesType on all previous layers
+        .map((layer) => ({
+          ...layer,
+          seriesType,
+        }))
     : [];
 
   const state: State = {
     legend: currentState ? currentState.legend : { isVisible: true, position: Position.Right },
     fittingFunction: currentState?.fittingFunction || 'None',
     preferredSeriesType: seriesType,
-    layers: [...keptLayers, newLayer],
+    layers: Object.keys(existingLayer).length ? keptLayers : [...keptLayers, newLayer],
   };
 
   return {

x-pack/test/functional/apps/lens/smokescreen.ts

@@ -165,5 +165,31 @@ export default function ({ getService, getPageObjects }: FtrProviderContext) {
       // legend item(s), so we're using a class selector here.
       expect(await find.allByCssSelector('.echLegendItem')).to.have.length(3);
     });
+
+    it('should switch from a multi-layer stacked bar to a multi-layer line chart', async () => {
+      await PageObjects.visualize.navigateToNewVisualization();
+      await PageObjects.visualize.clickVisType('lens');
+      await PageObjects.lens.goToTimeRange();
+
+      await PageObjects.lens.configureDimension({
+        dimension: 'lnsXY_xDimensionPanel > lns-empty-dimension',
+        operation: 'date_histogram',
+        field: '@timestamp',
+      });
+
+      await PageObjects.lens.configureDimension({
+        dimension: 'lnsXY_yDimensionPanel > lns-empty-dimension',
+        operation: 'avg',
+        field: 'bytes',
+      });
+
+      await PageObjects.lens.createLayer();
+
+      expect(await PageObjects.lens.hasChartSwitchWarning('line')).to.eql(false);
+
+      await PageObjects.lens.switchToVisualization('line');
+
+      expect(await PageObjects.lens.getLayerCount()).to.eql(2);
+    });
   });
 }

x-pack/test/functional/page_objects/lens_page.ts

@@ -167,5 +167,33 @@ export function LensPageProvider({ getService, getPageObjects }: FtrProviderCont
         await testSubjects.existOrFail('visTypeTitle');
       });
     },
+
+    /**
+     * Checks a specific subvisualization in the chart switcher for a "data loss" indicator
+     *
+     * @param subVisualizationId - the ID of the sub-visualization to switch to, such as
+     * lnsDatatable or bar_stacked
+     */
+    async hasChartSwitchWarning(subVisualizationId: string) {
+      await this.openChartSwitchPopover();
+
+      const element = await testSubjects.find(`lnsChartSwitchPopover_${subVisualizationId}`);
+      return await testSubjects.descendantExists('euiKeyPadMenuItem__betaBadgeWrapper', element);
+    },
+
+    /**
+     * Returns the number of layers visible in the chart configuration
+     */
+    async getLayerCount() {
+      const elements = await testSubjects.findAll('lnsLayerRemove');
+      return elements.length;
+    },
+
+    /**
+     * Adds a new layer to the chart, fails if the chart does not support new layers
+     */
+    async createLayer() {
+      await testSubjects.click('lnsLayerAddButton');
+    },
   });
 }