[Lens] Metric trendlines design changes

x-pack/plugins/lens/public/datasources/form_based/form_based.test.ts

@@ -1665,12 +1665,74 @@ describe('IndexPattern Data Source', () => {
         currentIndexPatternId: '1',
       };
       expect(FormBasedDatasource.removeLayer(state, 'first')).toEqual({
-        ...state,
+        removedLayerIds: ['first'],
+        newState: {
+          ...state,
+          layers: {
+            second: {
+              indexPatternId: '2',
+              columnOrder: [],
+              columns: {},
+            },
+          },
+        },
+      });
+    });
+
+    it('should remove linked layers', () => {
+      const state = {
         layers: {
+          first: {
+            indexPatternId: '1',
+            columnOrder: [],
+            columns: {},
+          },
           second: {
             indexPatternId: '2',
             columnOrder: [],
             columns: {},
+            linkToLayers: ['first'],
+          },
+        },
+        currentIndexPatternId: '1',
+      };
+      expect(FormBasedDatasource.removeLayer(state, 'first')).toEqual({
+        removedLayerIds: ['first', 'second'],
+        newState: {
+          ...state,
+          layers: {},
+        },
+      });
+    });
+  });
+
+  describe('#clearLayer', () => {
+    it('should clear a layer', () => {
+      const state = {
+        layers: {
+          first: {
+            indexPatternId: '1',
+            columnOrder: ['some', 'order'],
+            columns: {
+              some: {} as GenericIndexPatternColumn,
+              columns: {} as GenericIndexPatternColumn,
+            },
+            linkToLayers: ['some-layer'],
+          },
+        },
+        currentIndexPatternId: '1',
+      };
+      expect(FormBasedDatasource.clearLayer(state, 'first')).toEqual({
+        removedLayerIds: [],
+        newState: {
+          ...state,
+          layers: {
+            first: {
+              indexPatternId: '1',
+              columnOrder: [],
+              columns: {},
+              linkToLayers: ['some-layer'],
+            },
           },
         },
       });
@@ -1693,9 +1755,14 @@ describe('IndexPattern Data Source', () => {
         },
         currentIndexPatternId: '1',
       };
-      expect(FormBasedDatasource.removeLayer(state, 'first')).toEqual({
-        ...state,
-        layers: {},
+      expect(FormBasedDatasource.clearLayer(state, 'first')).toEqual({
+        removedLayerIds: ['second'],
+        newState: {
+          ...state,
+          layers: {
+            first: state.layers.first,
+          },
+        },
       });
     });
   });

x-pack/plugins/lens/public/datasources/form_based/form_based.tsx

@@ -218,27 +218,47 @@ export function getFormBasedDatasource({
     removeLayer(state: FormBasedPrivateState, layerId: string) {
       const newLayers = { ...state.layers };
       delete newLayers[layerId];
+      const removedLayerIds: string[] = [layerId];
 
       // delete layers linked to this layer
       Object.keys(newLayers).forEach((id) => {
         const linkedLayers = newLayers[id]?.linkToLayers;
         if (linkedLayers && linkedLayers.includes(layerId)) {
           delete newLayers[id];
+          removedLayerIds.push(id);
         }
       });
 
       return {
-        ...state,
-        layers: newLayers,
+        removedLayerIds,
+        newState: {
+          ...state,
+          layers: newLayers,
+        },
       };
     },
 
     clearLayer(state: FormBasedPrivateState, layerId: string) {
+      const newLayers = { ...state.layers };
+
+      const removedLayerIds: string[] = [];
+      // delete layers linked to this layer
+      Object.keys(newLayers).forEach((id) => {
+        const linkedLayers = newLayers[id]?.linkToLayers;
+        if (linkedLayers && linkedLayers.includes(layerId)) {
+          delete newLayers[id];
+          removedLayerIds.push(id);
+        }
+      });
+
       return {
-        ...state,
-        layers: {
-          ...state.layers,
-          [layerId]: blankLayer(state.currentIndexPatternId, state.layers[layerId].linkToLayers),
+        removedLayerIds,
+        newState: {
+          ...state,
+          layers: {
+            ...newLayers,
+            [layerId]: blankLayer(state.currentIndexPatternId, state.layers[layerId].linkToLayers),
+          },
         },
       };
     },

x-pack/plugins/lens/public/datasources/text_based/text_based_languages.test.ts

@@ -217,21 +217,24 @@ describe('IndexPattern Data Source', () => {
   describe('#removeLayer', () => {
     it('should remove a layer', () => {
       expect(TextBasedDatasource.removeLayer(baseState, 'a')).toEqual({
-        ...baseState,
-        layers: {
-          a: {
-            columns: [],
-            allColumns: [
-              {
-                columnId: 'col1',
-                fieldName: 'Test 1',
-                meta: {
-                  type: 'number',
+        removedLayerIds: ['a'],
+        newState: {
+          ...baseState,
+          layers: {
+            a: {
+              columns: [],
+              allColumns: [
+                {
+                  columnId: 'col1',
+                  fieldName: 'Test 1',
+                  meta: {
+                    type: 'number',
+                  },
                 },
-              },
-            ],
-            query: { sql: 'SELECT * FROM foo' },
-            index: 'foo',
+              ],
+              query: { sql: 'SELECT * FROM foo' },
+              index: 'foo',
+            },
           },
         },
       });

x-pack/plugins/lens/public/datasources/text_based/text_based_languages.tsx

@@ -281,18 +281,24 @@ export function getTextBasedDatasource({
       };
 
       return {
-        ...state,
-        layers: newLayers,
-        fieldList: state.fieldList,
+        removedLayerIds: [layerId],
+        newState: {
+          ...state,
+          layers: newLayers,
+          fieldList: state.fieldList,
+        },
       };
     },
 
     clearLayer(state: TextBasedPrivateState, layerId: string) {
       return {
-        ...state,
-        layers: {
-          ...state.layers,
-          [layerId]: { ...state.layers[layerId], columns: [] },
+        removedLayerIds: [],
+        newState: {
+          ...state,
+          layers: {
+            ...state.layers,
+            [layerId]: { ...state.layers[layerId], columns: [] },
+          },
         },
       };
     },

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

@@ -494,8 +494,8 @@ describe('chart_switch', () => {
 
     switchTo('visB', instance);
     expect(datasourceMap.testDatasource.removeLayer).toHaveBeenCalledWith({}, 'a');
-    expect(datasourceMap.testDatasource.removeLayer).toHaveBeenCalledWith(undefined, 'b');
-    expect(datasourceMap.testDatasource.removeLayer).toHaveBeenCalledWith(undefined, 'c');
+    expect(datasourceMap.testDatasource.removeLayer).toHaveBeenCalledWith({}, 'b');
+    expect(datasourceMap.testDatasource.removeLayer).toHaveBeenCalledWith({}, 'c');
     expect(visualizationMap.visB.getSuggestions).toHaveBeenCalledWith(
       expect.objectContaining({
         keptLayerIds: ['a'],

x-pack/plugins/lens/public/mocks/datasource_mock.ts

@@ -26,7 +26,7 @@ export function createMockDatasource(id: string): DatasourceMock {
 
   return {
     id: 'testDatasource',
-    clearLayer: jest.fn((state, _layerId) => state),
+    clearLayer: jest.fn((state, _layerId) => ({ newState: state, removedLayerIds: [] })),
     getDatasourceSuggestionsForField: jest.fn((_state, _item, filterFn, _indexPatterns) => []),
     getDatasourceSuggestionsForVisualizeField: jest.fn(
       (_state, _indexpatternId, _fieldName, _indexPatterns) => []
@@ -44,7 +44,7 @@ export function createMockDatasource(id: string): DatasourceMock {
     renderLayerPanel: jest.fn(),
     toExpression: jest.fn((_frame, _state, _indexPatterns) => null),
     insertLayer: jest.fn((_state, _newLayerId) => ({})),
-    removeLayer: jest.fn((_state, _layerId) => {}),
+    removeLayer: jest.fn((state, layerId) => ({ newState: state, removedLayerIds: [layerId] })),
     cloneLayer: jest.fn((_state, _layerId, _newLayerId, getNewId) => {}),
     removeColumn: jest.fn((props) => {}),
     getLayers: jest.fn((_state) => []),

x-pack/plugins/lens/public/state_management/lens_slice.test.ts

@@ -39,6 +39,7 @@ describe('lensSlice', () => {
   let store: EnhancedStore<{ lens: LensAppState }>;
   beforeEach(() => {
     store = makeLensStore({}).store;
+    jest.clearAllMocks();
   });
   const customQuery = { query: 'custom' } as Query;
 
@@ -275,17 +276,21 @@ describe('lensSlice', () => {
         return {
           id: datasourceId,
           getPublicAPI: () => ({
-            datasourceId: 'testDatasource',
+            datasourceId,
             getOperationForColumnId: jest.fn(),
             getTableSpec: jest.fn(),
           }),
           getLayers: () => ['layer1'],
-          clearLayer: (layerIds: unknown, layerId: string) =>
-            (layerIds as string[]).map((id: string) =>
+          clearLayer: (layerIds: unknown, layerId: string) => ({
+            removedLayerIds: [],
+            newState: (layerIds as string[]).map((id: string) =>
               id === layerId ? `${datasourceId}_clear_${layerId}` : id
             ),
-          removeLayer: (layerIds: unknown, layerId: string) =>
-            (layerIds as string[]).filter((id: string) => id !== layerId),
+          }),
+          removeLayer: (layerIds: unknown, layerId: string) => ({
+            newState: (layerIds as string[]).filter((id: string) => id !== layerId),
+            removedLayerIds: [layerId],
+          }),
           insertLayer: (layerIds: unknown, layerId: string, layersToLinkTo: string[]) => [
             ...(layerIds as string[]),
             layerId,
@@ -317,8 +322,9 @@ describe('lensSlice', () => {
             (layerIds as string[]).map((id: string) =>
               id === layerId ? `vis_clear_${layerId}` : id
             ),
-          removeLayer: (layerIds: unknown, layerId: string) =>
-            (layerIds as string[]).filter((id: string) => id !== layerId),
+          removeLayer: jest.fn((layerIds: unknown, layerId: string) =>
+            (layerIds as string[]).filter((id: string) => id !== layerId)
+          ),
           getLayerIds: (layerIds: unknown) => layerIds as string[],
           getLayersToLinkTo: (state, newLayerId) => ['linked-layer-id'],
           appendLayer: (layerIds: unknown, layerId: string) => [...(layerIds as string[]), layerId],
@@ -482,6 +488,54 @@ describe('lensSlice', () => {
         expect(state.datasourceStates.testDatasource2.state).toEqual(['layer2']);
         expect(state.stagedPreview).not.toBeDefined();
       });
+
+      it('removeLayer: should remove all layers from visualization that were removed by datasource', () => {
+        const removedLayerId = 'other-removed-layer';
+
+        const testDatasource3 = testDatasource('testDatasource3');
+        testDatasource3.removeLayer = (layerIds: unknown, layerId: string) => ({
+          newState: (layerIds as string[]).filter((id: string) => id !== layerId),
+          removedLayerIds: [layerId, removedLayerId],
+        });
+
+        const localStore = makeLensStore({
+          preloadedState: {
+            activeDatasourceId: 'testDatasource',
+            datasourceStates: {
+              ...datasourceStates,
+              testDatasource3: {
+                isLoading: false,
+                state: [],
+              },
+            },
+            visualization: {
+              activeId: activeVisId,
+              state: ['layer1', 'layer2'],
+            },
+            stagedPreview: {
+              visualization: {
+                activeId: activeVisId,
+                state: ['layer1', 'layer2'],
+              },
+              datasourceStates,
+            },
+          },
+          storeDeps: mockStoreDeps({
+            visualizationMap: visualizationMap as unknown as VisualizationMap,
+            datasourceMap: { ...datasourceMap, testDatasource3 } as unknown as DatasourceMap,
+          }),
+        }).store;
+
+        localStore.dispatch(
+          removeOrClearLayer({
+            visualizationId: 'testVis',
+            layerId: 'layer1',
+            layerIds: ['layer1', 'layer2'],
+          })
+        );
+
+        expect(visualizationMap[activeVisId].removeLayer).toHaveBeenCalledTimes(2);
+      });
     });
 
     describe('removing a dimension', () => {
@@ -546,8 +600,6 @@ describe('lensSlice', () => {
             datasourceMap: datasourceMap as unknown as DatasourceMap,
           }),
         }).store;
-
-        jest.clearAllMocks();
       });
 
       it('removes a dimension', () => {