fix PR comment and add test

Signed-off-by: Anan Zhuang <[email protected]>

src/plugins/discover/public/application/utils/state_management/discover_slice.test.tsx

@@ -4,6 +4,7 @@
  */
 
 import { discoverSlice, DiscoverState } from './discover_slice';
+import { SortOrder } from '../../../saved_searches/types';
 
 describe('discoverSlice', () => {
   let initialState: DiscoverState;
@@ -134,4 +135,40 @@ describe('discoverSlice', () => {
     const result = discoverSlice.reducer(initialState, action);
     expect(result.columns).toEqual(['column1', 'column2', 'column3']);
   });
+
+  it('should set the savedQuery when a valid string is provided', () => {
+    const savedQueryId = 'some-query-id';
+    const action = { type: 'discover/setSavedQuery', payload: savedQueryId };
+    const result = discoverSlice.reducer(initialState, action);
+    expect(result.savedQuery).toEqual(savedQueryId);
+  });
+
+  it('should remove the savedQuery from state when payload is undefined', () => {
+    // pre-set the savedQuery in the initialState
+    const initialStateWithSavedQuery = {
+      ...initialState,
+      savedQuery: 'existing-query-id',
+    };
+
+    const action = { type: 'discover/setSavedQuery', payload: undefined };
+    const result = discoverSlice.reducer(initialStateWithSavedQuery, action);
+
+    // Check that savedQuery is not in the resulting state
+    expect(result.savedQuery).toBeUndefined();
+  });
+
+  it('should not affect other state properties when setting savedQuery', () => {
+    const initialStateWithOtherProperties = {
+      ...initialState,
+      columns: ['column1', 'column2'],
+      sort: [['field1', 'asc']] as SortOrder[],
+    };
+    const savedQueryId = 'new-query-id';
+    const action = { type: 'discover/setSavedQuery', payload: savedQueryId };
+    const result = discoverSlice.reducer(initialStateWithOtherProperties, action);
+    // check that other properties remain unchanged
+    expect(result.columns).toEqual(['column1', 'column2']);
+    expect(result.sort).toEqual([['field1', 'asc']] as SortOrder[]);
+    expect(result.savedQuery).toEqual(savedQueryId);
+  });
 });

src/plugins/discover/public/application/utils/state_management/discover_slice.tsx

@@ -192,11 +192,17 @@ export const discoverSlice = createSlice({
         },
       };
     },
-    setSavedQuery(state, action: PayloadAction<string>) {
-      return {
-        ...state,
-        savedQuery: action.payload,
-      };
+    setSavedQuery(state, action: PayloadAction<string | undefined>) {
+      if (action.payload === undefined) {
+        // if the payload is undefined, remove the savedQuery property
+        const { savedQuery, ...restState } = state;
+        return restState;
+      } else {
+        return {
+          ...state,
+          savedQuery: action.payload,
+        };
+      }
     },
   },
 });

src/plugins/discover/public/application/view_components/canvas/top_nav.tsx

@@ -14,7 +14,7 @@ import { getTopNavLinks } from '../../components/top_nav/get_top_nav_links';
 import { useDiscoverContext } from '../context';
 import { getRootBreadcrumbs } from '../../helpers/breadcrumbs';
 import { opensearchFilters, connectStorageToQueryState } from '../../../../../data/public';
-import { useDispatch, setSavedQuery, useSelector, setState } from '../../utils/state_management';
+import { useDispatch, setSavedQuery, useSelector } from '../../utils/state_management';
 
 export interface TopNavProps {
   opts: {
@@ -78,14 +78,7 @@ export const TopNav = ({ opts, showSaveQuery }: TopNavProps) => {
   ]);
 
   const updateSavedQueryId = (newSavedQueryId: string | undefined) => {
-    if (newSavedQueryId) {
-      dispatch(setSavedQuery(newSavedQueryId));
-    } else {
-      // remove savedQueryId from state
-      const newState = { ...state };
-      delete newState.savedQuery;
-      dispatch(setState(newState));
-    }
+    dispatch(setSavedQuery(newSavedQueryId));
   };
 
   return (

src/plugins/vis_builder/public/application/components/top_nav.tsx

@@ -13,12 +13,13 @@ import { VisBuilderServices } from '../../types';
 import './top_nav.scss';
 import { useIndexPatterns, useSavedVisBuilderVis } from '../utils/use';
 import { useTypedSelector, useTypedDispatch } from '../utils/state_management';
-import { setSavedQuery, setState } from '../utils/state_management/visualization_slice';
+import { setSavedQuery } from '../utils/state_management/visualization_slice';
 import { setEditorState } from '../utils/state_management/metadata_slice';
 import { useCanSave } from '../utils/use/use_can_save';
 import { saveStateToSavedObject } from '../../saved_visualizations/transforms';
 import { TopNavMenuData } from '../../../../navigation/public';
 import { opensearchFilters, connectStorageToQueryState } from '../../../../data/public';
+import { RootState } from '../../../../data_explorer/public';
 
 export const TopNav = () => {
   // id will only be set for the edit route
@@ -32,7 +33,7 @@ export const TopNav = () => {
     appName,
     capabilities,
   } = services;
-  const rootState = useTypedSelector((state) => state);
+  const rootState = useTypedSelector((state: RootState) => state);
   const dispatch = useTypedDispatch();
 
   const saveDisabledReason = useCanSave();
@@ -81,14 +82,7 @@ export const TopNav = () => {
   });
 
   const updateSavedQueryId = (newSavedQueryId: string | undefined) => {
-    if (newSavedQueryId) {
-      dispatch(setSavedQuery(newSavedQueryId));
-    } else {
-      // remove savedQueryId from state
-      const newState = rootState;
-      delete newState.visualization.savedQuery;
-      dispatch(setState(newState.visualization));
-    }
+    dispatch(setSavedQuery(newSavedQueryId));
   };
   const showSaveQuery = !!capabilities['visualization-visbuilder']?.saveQuery;
 

src/plugins/vis_builder/public/application/utils/state_management/visualization_slice.test.ts

@@ -0,0 +1,148 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import {
+  slice as visualizationSlice,
+  VisualizationState,
+  setIndexPattern,
+  setSearchField,
+  editDraftAgg,
+  saveDraftAgg,
+  updateAggConfigParams,
+  setAggParamValue,
+  reorderAgg,
+  setSavedQuery,
+  setState,
+} from './visualization_slice';
+import { CreateAggConfigParams } from '../../../../../data/common';
+
+describe('visualizationSlice', () => {
+  let initialState: VisualizationState;
+
+  beforeEach(() => {
+    initialState = {
+      searchField: '',
+      activeVisualization: {
+        name: 'some-vis',
+        aggConfigParams: [],
+      },
+      savedQuery: undefined,
+    } as VisualizationState;
+  });
+
+  it('should handle setIndexPattern', () => {
+    const indexPattern = 'new-index-pattern';
+    const action = setIndexPattern(indexPattern);
+    const result = visualizationSlice.reducer(initialState, action);
+    expect(result.indexPattern).toEqual(indexPattern);
+  });
+
+  it('should handle setSearchField', () => {
+    const searchField = 'new-search-field';
+    const action = setSearchField(searchField);
+    const result = visualizationSlice.reducer(initialState, action);
+    expect(result.searchField).toEqual(searchField);
+  });
+
+  it('should handle editDraftAgg', () => {
+    const draftAgg: CreateAggConfigParams = { id: '1', enabled: true, type: 'count', params: {} };
+    const action = editDraftAgg(draftAgg);
+    const result = visualizationSlice.reducer(initialState, action);
+    expect(result.activeVisualization?.draftAgg).toEqual(draftAgg);
+  });
+
+  it('should handle saveDraftAgg', () => {
+    const draftAgg: CreateAggConfigParams = { id: '1', enabled: true, type: 'count', params: {} };
+    const stateWithDraft = {
+      ...initialState,
+      activeVisualization: {
+        name: 'some-name',
+        draftAgg,
+        aggConfigParams: [],
+      },
+    };
+    const action = saveDraftAgg(undefined);
+    const result = visualizationSlice.reducer(stateWithDraft, action);
+    expect(result.activeVisualization?.aggConfigParams).toContain(draftAgg);
+  });
+
+  it('should handle updateAggConfigParams', () => {
+    const newAggConfigParams: CreateAggConfigParams[] = [
+      { id: '2', enabled: true, type: 'avg', params: {} },
+    ];
+    const action = updateAggConfigParams(newAggConfigParams);
+    const result = visualizationSlice.reducer(initialState, action);
+    expect(result.activeVisualization?.aggConfigParams).toEqual(newAggConfigParams);
+  });
+
+  it('should handle setAggParamValue', () => {
+    const aggParamValue = { aggId: '1', paramName: 'field', value: 'newField' };
+    const action = setAggParamValue(aggParamValue);
+
+    const initialStateWithAgg = {
+      ...initialState,
+      activeVisualization: {
+        ...initialState.activeVisualization,
+        name: 'defaultName',
+        aggConfigParams: [{ id: '1', enabled: true, type: 'count', params: {} }],
+      },
+    };
+
+    const result = visualizationSlice.reducer(initialStateWithAgg, action);
+    expect(
+      result.activeVisualization?.aggConfigParams?.[0]?.params?.[aggParamValue.paramName]
+    ).toEqual(aggParamValue.value);
+  });
+
+  it('should handle reorderAgg', () => {
+    const reorderAction = { sourceId: '1', destinationId: '2' };
+    const action = reorderAgg(reorderAction);
+
+    // Initial state with multiple aggregations
+    const initialStateWithMultipleAggs = {
+      ...initialState,
+      activeVisualization: {
+        ...initialState.activeVisualization,
+        name: 'defaultName',
+        aggConfigParams: [
+          { id: '1', enabled: true, type: 'count', params: {} },
+          { id: '2', enabled: true, type: 'avg', params: {} },
+        ],
+      },
+    };
+
+    const result = visualizationSlice.reducer(initialStateWithMultipleAggs, action);
+    // verify the order of aggConfigParams
+    expect(result.activeVisualization?.aggConfigParams[0].id).toEqual('2');
+    expect(result.activeVisualization?.aggConfigParams[1].id).toEqual('1');
+  });
+
+  it('should handle savedQueryId by setSavedQuery', () => {
+    const savedQueryId = 'some-query-id';
+    const action = setSavedQuery(savedQueryId);
+    const result = visualizationSlice.reducer(initialState, action) as VisualizationState;
+    expect(result.savedQuery).toEqual(savedQueryId);
+  });
+
+  it('should handle undefined savedQueryId by setSavedQuery', () => {
+    const savedQueryId = undefined;
+    const action = setSavedQuery(savedQueryId);
+    const result = visualizationSlice.reducer(initialState, action) as VisualizationState;
+    expect(result.savedQuery).toBeUndefined();
+  });
+
+  it('should handle setState', () => {
+    const newState = {
+      searchField: 'new-search-field',
+      activeVisualization: {
+        name: 'new-vis',
+        aggConfigParams: [],
+      },
+    };
+    const action = setState(newState);
+    const result = visualizationSlice.reducer(initialState, action);
+    expect(result).toEqual(newState);
+  });
+});

src/plugins/vis_builder/public/application/utils/state_management/visualization_slice.ts

@@ -116,11 +116,17 @@ export const slice = createSlice({
         [action.payload.paramName]: action.payload.value,
       };
     },
-    setSavedQuery(state, action: PayloadAction<string>) {
-      return {
-        ...state,
-        savedQuery: action.payload,
-      };
+    setSavedQuery(state, action: PayloadAction<string | undefined>) {
+      if (action.payload === undefined) {
+        // if the payload is undefined, remove the savedQuery property
+        const { savedQuery, ...restState } = state;
+        return restState;
+      } else {
+        return {
+          ...state,
+          savedQuery: action.payload,
+        };
+      }
     },
     setState: (_state, action: PayloadAction<VisualizationState>) => {
       return action.payload;