From 948b99ef03f434db806a4d8f1707cfc8ee11b85a Mon Sep 17 00:00:00 2001
From: Marta Bondyra <marta.bondyra@elastic.co>
Date: Tue, 15 Oct 2024 14:11:38 +0200
Subject: [PATCH] [Lens] fix suggestions api (allow not pass layerId to switch
 VisualizationType)

---
 .../visualizations/xy/visualization.test.tsx  | 26 +++++++++++++++++++
 .../visualizations/xy/visualization.tsx       |  3 +--
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/x-pack/plugins/lens/public/visualizations/xy/visualization.test.tsx b/x-pack/plugins/lens/public/visualizations/xy/visualization.test.tsx
index 012823831b8eb..e2c6ce25bd2e3 100644
--- a/x-pack/plugins/lens/public/visualizations/xy/visualization.test.tsx
+++ b/x-pack/plugins/lens/public/visualizations/xy/visualization.test.tsx
@@ -4214,4 +4214,30 @@ describe('xy_visualization', () => {
       `);
     });
   });
+  describe('switchVisualizationType', () => {
+    it('should switch all the layers to the new visualization type if layerId is not specified (AI assistant case)', () => {
+      const state = exampleState();
+      state.layers[1] = state.layers[0];
+      state.layers[1].layerId = 'second';
+      state.layers[2] = state.layers[0];
+      state.layers[2].layerId = 'third';
+      const newType = 'bar';
+      const newState = xyVisualization.switchVisualizationType!(newType, state);
+      expect((newState.layers[0] as XYDataLayerConfig).seriesType).toEqual(newType);
+      expect((newState.layers[1] as XYDataLayerConfig).seriesType).toEqual(newType);
+      expect((newState.layers[2] as XYDataLayerConfig).seriesType).toEqual(newType);
+    });
+    it('should switch only the second layer to the new visualization type if layerId is specified (chart switch case)', () => {
+      const state = exampleState();
+      state.layers[1] = { ...state.layers[0] };
+      state.layers[1].layerId = 'second';
+      state.layers[2] = { ...state.layers[0] };
+      state.layers[2].layerId = 'third';
+      const newType = 'bar';
+      const newState = xyVisualization.switchVisualizationType!(newType, state, 'first');
+      expect((newState.layers[0] as XYDataLayerConfig).seriesType).toEqual(newType);
+      expect((newState.layers[1] as XYDataLayerConfig).seriesType).toEqual('area');
+      expect((newState.layers[2] as XYDataLayerConfig).seriesType).toEqual('area');
+    });
+  });
 });
diff --git a/x-pack/plugins/lens/public/visualizations/xy/visualization.tsx b/x-pack/plugins/lens/public/visualizations/xy/visualization.tsx
index 64a2ad4fc2754..171d2a9b8b185 100644
--- a/x-pack/plugins/lens/public/visualizations/xy/visualization.tsx
+++ b/x-pack/plugins/lens/public/visualizations/xy/visualization.tsx
@@ -263,14 +263,13 @@ export const getXyVisualization = ({
   getDescription,
 
   switchVisualizationType(seriesType: string, state: State, layerId?: string) {
-    const dataLayer = state.layers.find((l) => l.layerId === layerId);
+    const dataLayer = state.layers.find((l) => l.layerId === layerId) ?? state.layers[0];
     if (dataLayer && !isDataLayer(dataLayer)) {
       throw new Error('Cannot switch series type for non-data layer');
     }
     if (!dataLayer) {
       return state;
     }
-    // todo: test how they switch between percentage etc
     const currentStackingType = stackingTypes.find(({ subtypes }) =>
       subtypes.includes(dataLayer.seriesType)
     );