diff --git a/tensorboard/webapp/feature_flag/actions/feature_flag_actions.ts b/tensorboard/webapp/feature_flag/actions/feature_flag_actions.ts index a7fe4e2cd5..41673db0a2 100644 --- a/tensorboard/webapp/feature_flag/actions/feature_flag_actions.ts +++ b/tensorboard/webapp/feature_flag/actions/feature_flag_actions.ts @@ -43,3 +43,14 @@ export const featureFlagOverrideChanged = createAction( flags: Partial; }>() ); + +export const featureFlagOverridesReset = createAction( + '[FEATURE FLAG] Reset feature flag overrides', + props<{ + flags: Array; + }>() +); + +export const allFeatureFlagOverridesReset = createAction( + '[FEATURE FLAG] Reset all feature flag overrides' +); diff --git a/tensorboard/webapp/feature_flag/effects/BUILD b/tensorboard/webapp/feature_flag/effects/BUILD index 628544adc2..2abce8bab0 100644 --- a/tensorboard/webapp/feature_flag/effects/BUILD +++ b/tensorboard/webapp/feature_flag/effects/BUILD @@ -10,6 +10,7 @@ tf_ng_module( deps = [ "//tensorboard/webapp:tb_polymer_interop_types", "//tensorboard/webapp/feature_flag:force_svg_data_source", + "//tensorboard/webapp/feature_flag:types", "//tensorboard/webapp/feature_flag/actions", "//tensorboard/webapp/feature_flag/store", "//tensorboard/webapp/feature_flag/store:types", diff --git a/tensorboard/webapp/feature_flag/effects/feature_flag_effects.ts b/tensorboard/webapp/feature_flag/effects/feature_flag_effects.ts index 5d7fe09c3f..2dcdb47090 100644 --- a/tensorboard/webapp/feature_flag/effects/feature_flag_effects.ts +++ b/tensorboard/webapp/feature_flag/effects/feature_flag_effects.ts @@ -20,7 +20,9 @@ import {combineLatestWith, map, tap, withLatestFrom} from 'rxjs/operators'; import '../../tb_polymer_interop_types'; import {TBFeatureFlagDataSource} from '../../webapp_data_source/tb_feature_flag_data_source_types'; import { + allFeatureFlagOverridesReset, featureFlagOverrideChanged, + featureFlagOverridesReset, partialFeatureFlagsLoaded, } from '../actions/feature_flag_actions'; import {ForceSvgDataSource} from '../force_svg_data_source'; @@ -95,6 +97,30 @@ export class FeatureFlagEffects { {dispatch: false} ); + readonly resetFeatureFlagOverrides$ = createEffect( + () => + this.actions$.pipe( + ofType(featureFlagOverridesReset), + tap(({flags}) => { + flags.forEach((flag) => { + this.dataSource.resetPersistedFeatureFlag(flag); + }); + }) + ), + {dispatch: false} + ); + + readonly resetAllFeatureFlagOverrides$ = createEffect( + () => + this.actions$.pipe( + ofType(allFeatureFlagOverridesReset), + tap(() => { + this.dataSource.resetAllPersistedFeatureFlags(); + }) + ), + {dispatch: false} + ); + constructor( private readonly actions$: Actions, private readonly store: Store, diff --git a/tensorboard/webapp/feature_flag/effects/feature_flag_effects_test.ts b/tensorboard/webapp/feature_flag/effects/feature_flag_effects_test.ts index 4c5bc2000c..e58ef05e8f 100644 --- a/tensorboard/webapp/feature_flag/effects/feature_flag_effects_test.ts +++ b/tensorboard/webapp/feature_flag/effects/feature_flag_effects_test.ts @@ -23,7 +23,9 @@ import { TestingTBFeatureFlagDataSource, } from '../../webapp_data_source/tb_feature_flag_testing'; import { + allFeatureFlagOverridesReset, featureFlagOverrideChanged, + featureFlagOverridesReset, partialFeatureFlagsLoaded, } from '../actions/feature_flag_actions'; import {ForceSvgDataSource} from '../force_svg_data_source'; @@ -200,4 +202,28 @@ describe('feature_flag_effects', () => { }); }); }); + + describe('resetFeatureFlagOverrides', () => { + it('calls resetPersistedFeatureFlag', () => { + const resetFlagSpy = spyOn( + dataSource, + 'resetPersistedFeatureFlag' + ).and.stub(); + effects.resetFeatureFlagOverrides$.subscribe(); + actions.next(featureFlagOverridesReset({flags: ['inColab']})); + expect(resetFlagSpy).toHaveBeenCalledOnceWith('inColab'); + }); + }); + + describe('resetAllFeatureFlagOverrides', () => { + it('calls resetAllPersistedFeatureFlags', () => { + const resetAllFlagsSpy = spyOn( + dataSource, + 'resetAllPersistedFeatureFlags' + ).and.stub(); + effects.resetAllFeatureFlagOverrides$.subscribe(); + actions.next(allFeatureFlagOverridesReset()); + expect(resetAllFlagsSpy).toHaveBeenCalledOnceWith(); + }); + }); }); diff --git a/tensorboard/webapp/feature_flag/store/feature_flag_reducers.ts b/tensorboard/webapp/feature_flag/store/feature_flag_reducers.ts index 7426381eee..a014f52c48 100644 --- a/tensorboard/webapp/feature_flag/store/feature_flag_reducers.ts +++ b/tensorboard/webapp/feature_flag/store/feature_flag_reducers.ts @@ -15,6 +15,7 @@ limitations under the License. import {Action, createReducer, on} from '@ngrx/store'; import {globalSettingsLoaded, ThemeValue} from '../../persistent_settings'; import * as actions from '../actions/feature_flag_actions'; +import {FeatureFlags} from '../types'; import {initialState} from './feature_flag_store_config_provider'; import {FeatureFlagState} from './feature_flag_types'; @@ -43,6 +44,34 @@ const reducer = createReducer( }, }; }), + on(actions.featureFlagOverrideChanged, (state, newOverrides) => { + return { + ...state, + flagOverrides: { + ...state.flagOverrides, + ...newOverrides.flags, + }, + }; + }), + on(actions.featureFlagOverridesReset, (state, overrides) => { + if (!overrides || !overrides.flags || !overrides.flags.length) { + return state; + } + const flagOverrides = {...state.flagOverrides}; + overrides.flags.forEach((key) => { + delete flagOverrides[key as keyof FeatureFlags]; + }); + return { + ...state, + flagOverrides, + }; + }), + on(actions.allFeatureFlagOverridesReset, (state) => { + return { + ...state, + flagOverrides: {}, + }; + }), on(globalSettingsLoaded, (state, {partialSettings}) => { if (!partialSettings.themeOverride) { return state; diff --git a/tensorboard/webapp/feature_flag/store/feature_flag_reducers_test.ts b/tensorboard/webapp/feature_flag/store/feature_flag_reducers_test.ts index daeeafdb29..4cb5804ec8 100644 --- a/tensorboard/webapp/feature_flag/store/feature_flag_reducers_test.ts +++ b/tensorboard/webapp/feature_flag/store/feature_flag_reducers_test.ts @@ -108,6 +108,119 @@ describe('feature_flag_reducers', () => { }); }); + describe('#featureFlagOverrideChanged', () => { + it('does not remove existing overrides', () => { + const prevState = buildFeatureFlagState({ + flagOverrides: { + inColab: true, + defaultEnableDarkMode: true, + }, + }); + const nextState = reducers( + prevState, + actions.featureFlagOverrideChanged({ + flags: {enabledLinkedTime: true}, + }) + ); + expect(nextState.flagOverrides).toEqual({ + inColab: true, + defaultEnableDarkMode: true, + enabledLinkedTime: true, + }); + }); + + it('changes values of existing overrides', () => { + const prevState = buildFeatureFlagState({ + flagOverrides: { + inColab: true, + defaultEnableDarkMode: true, + }, + }); + const nextState = reducers( + prevState, + actions.featureFlagOverrideChanged({ + flags: {inColab: false}, + }) + ); + expect(nextState.flagOverrides).toEqual({ + inColab: false, + defaultEnableDarkMode: true, + }); + }); + }); + + describe('#featureFlagOverridesReset', () => { + it('does nothing when no overrides are provided', () => { + const prevState = buildFeatureFlagState(); + const nextState = reducers( + prevState, + actions.featureFlagOverridesReset({flags: []}) + ); + // Intentionally using toBe rather than toEqual to ensure the object is the SAME object (not a copy). + expect(nextState).toBe(prevState); + }); + + it('removes all provided overrides', () => { + const prevState = buildFeatureFlagState({ + flagOverrides: { + inColab: true, + forceSvg: true, + scalarsBatchSize: 5, + }, + }); + const nextState = reducers( + prevState, + actions.featureFlagOverridesReset({flags: ['forceSvg', 'inColab']}) + ); + expect(nextState.flagOverrides).toEqual({ + scalarsBatchSize: 5, + }); + }); + + it('does not effect default flags', () => { + const prevState = buildFeatureFlagState({ + flagOverrides: { + inColab: true, + }, + }); + + prevState.defaultFlags.inColab = true; + const nextState = reducers( + prevState, + actions.featureFlagOverridesReset({flags: ['inColab']}) + ); + expect(nextState.flagOverrides).toEqual({}); + expect(nextState.defaultFlags).toBe(prevState.defaultFlags); + }); + }); + + describe('#allFeatureFlagOverridesReset', () => { + it('always generates a new state', () => { + const prevState = buildFeatureFlagState({flagOverrides: {}}); + const nextState = reducers( + prevState, + actions.allFeatureFlagOverridesReset() + ); + expect(nextState.flagOverrides).not.toBe(prevState.flagOverrides); + expect(nextState.flagOverrides).toEqual({}); + }); + + it('removes all overridden feature flags', () => { + const prevState = buildFeatureFlagState({ + flagOverrides: { + inColab: true, + forceSvg: true, + scalarsBatchSize: 5, + }, + }); + const nextState = reducers( + prevState, + actions.allFeatureFlagOverridesReset() + ); + expect(nextState.flagOverrides).toEqual({}); + }); + }); + describe('#globalSettingsLoaded', () => { it('sets dark mode overrides when global settings include it', () => { const prevState = buildFeatureFlagState({ diff --git a/tensorboard/webapp/feature_flag/store/feature_flag_selectors.ts b/tensorboard/webapp/feature_flag/store/feature_flag_selectors.ts index 41052d8399..e456bfa9bd 100644 --- a/tensorboard/webapp/feature_flag/store/feature_flag_selectors.ts +++ b/tensorboard/webapp/feature_flag/store/feature_flag_selectors.ts @@ -42,6 +42,13 @@ export const getFeatureFlags = createSelector( } ); +export const getDefaultFeatureFlags = createSelector( + selectFeatureFlagState, + (state: FeatureFlagState): FeatureFlags => { + return state.defaultFlags; + } +); + export const getOverriddenFeatureFlags = createSelector( selectFeatureFlagState, (state: FeatureFlagState): Partial => { diff --git a/tensorboard/webapp/feature_flag/store/feature_flag_selectors_test.ts b/tensorboard/webapp/feature_flag/store/feature_flag_selectors_test.ts index b98b880c2a..2eaed024e2 100644 --- a/tensorboard/webapp/feature_flag/store/feature_flag_selectors_test.ts +++ b/tensorboard/webapp/feature_flag/store/feature_flag_selectors_test.ts @@ -55,6 +55,29 @@ describe('feature_flag_selectors', () => { }); }); + describe('#getDefaultFeatureFlags', () => { + it('returns exact copy of the default feature flags', () => { + const initialState = buildFeatureFlagState(); + const state = buildState(initialState); + const actual = selectors.getDefaultFeatureFlags(state); + + expect(actual).toEqual(initialState.defaultFlags); + }); + + it('is uneffected by feature flag overrides', () => { + const initialState = buildFeatureFlagState({ + flagOverrides: { + forceSvg: true, + inColab: true, + }, + }); + const state = buildState(initialState); + const actual = selectors.getDefaultFeatureFlags(state); + + expect(actual).toEqual(initialState.defaultFlags); + }); + }); + describe('#getOverriddenFeatureFlags', () => { it('returns empty object if it is not overridden', () => { const state = buildState(buildFeatureFlagState()); diff --git a/tensorboard/webapp/webapp_data_source/tb_feature_flag_data_source.ts b/tensorboard/webapp/webapp_data_source/tb_feature_flag_data_source.ts index 525bd4a03f..9d28b134ad 100644 --- a/tensorboard/webapp/webapp_data_source/tb_feature_flag_data_source.ts +++ b/tensorboard/webapp/webapp_data_source/tb_feature_flag_data_source.ts @@ -74,6 +74,10 @@ export class FeatureFlagOverrideDataSource implements TBFeatureFlagDataSource { ); } + resetAllPersistedFeatureFlags() { + localStorage.removeItem(FEATURE_FLAG_STORAGE_KEY); + } + getPersistentFeatureFlags(): Partial { const currentState = localStorage.getItem(FEATURE_FLAG_STORAGE_KEY); if (currentState == null) { diff --git a/tensorboard/webapp/webapp_data_source/tb_feature_flag_data_source_test.ts b/tensorboard/webapp/webapp_data_source/tb_feature_flag_data_source_test.ts index 7e7f443e7d..164d84946e 100644 --- a/tensorboard/webapp/webapp_data_source/tb_feature_flag_data_source_test.ts +++ b/tensorboard/webapp/webapp_data_source/tb_feature_flag_data_source_test.ts @@ -349,5 +349,20 @@ describe('tb_feature_flag_data_source', () => { ); }); }); + + describe('resetAllPersistedFeatureFlags', () => { + it('removes entry from localStorage', () => { + spyOn(localStorage, 'getItem').and.returnValue( + '{"enabledScalarDataTable":true}' + ); + const setItemSpy = spyOn(localStorage, 'setItem').and.stub(); + const removeItemSpy = spyOn(localStorage, 'removeItem').and.stub(); + + dataSource.resetAllPersistedFeatureFlags(); + expect(removeItemSpy).toHaveBeenCalledOnceWith( + 'tb_feature_flag_storage_key' + ); + }); + }); }); }); diff --git a/tensorboard/webapp/webapp_data_source/tb_feature_flag_data_source_types.ts b/tensorboard/webapp/webapp_data_source/tb_feature_flag_data_source_types.ts index 696bb3c7e4..e4e2040ab9 100644 --- a/tensorboard/webapp/webapp_data_source/tb_feature_flag_data_source_types.ts +++ b/tensorboard/webapp/webapp_data_source/tb_feature_flag_data_source_types.ts @@ -48,6 +48,11 @@ export abstract class TBFeatureFlagDataSource { featureFlag: K ): void; + /** + * Removes all feature flags overridden in localStorage. + */ + abstract resetAllPersistedFeatureFlags(): void; + /** * Gets the serialized data stored in localStorage for the stored feature * flags. diff --git a/tensorboard/webapp/webapp_data_source/tb_feature_flag_testing.ts b/tensorboard/webapp/webapp_data_source/tb_feature_flag_testing.ts index cfad5c7032..f19879e7b1 100644 --- a/tensorboard/webapp/webapp_data_source/tb_feature_flag_testing.ts +++ b/tensorboard/webapp/webapp_data_source/tb_feature_flag_testing.ts @@ -27,6 +27,8 @@ export class TestingTBFeatureFlagDataSource extends TBFeatureFlagDataSource { resetPersistedFeatureFlag(featureFlag: K) {} + resetAllPersistedFeatureFlags(): void {} + getPersistentFeatureFlags(): Partial { return {}; }