diff --git a/packages/react/src/context/use-context-mutator.ts b/packages/react/src/context/use-context-mutator.ts index 800688baf..8ab11ac47 100644 --- a/packages/react/src/context/use-context-mutator.ts +++ b/packages/react/src/context/use-context-mutator.ts @@ -18,10 +18,10 @@ export type ContextMutation = { * Context-aware function to set the desired context (see: {@link ContextMutationOptions} for details). * There's generally no need to await the result of this function; flag evaluation hooks will re-render when the context is updated. * This promise never rejects. - * @param updatedContext + * @param updatedContext New context object or method to generate it from the current context * @returns Promise for awaiting the context update */ - setContext: (updatedContext: EvaluationContext) => Promise; + setContext: (updatedContext: EvaluationContext | ((currentContext: EvaluationContext) => EvaluationContext)) => Promise; }; /** @@ -34,16 +34,20 @@ export function useContextMutator(options: ContextMutationOptions = { defaultCon const { domain } = useContext(Context) || {}; const previousContext = useRef(null); - const setContext = useCallback(async (updatedContext: EvaluationContext) => { - if (previousContext.current !== updatedContext) { + const setContext = useCallback(async (updatedContext: EvaluationContext | ((currentContext: EvaluationContext) => EvaluationContext)): Promise => { + const resolvedContext = typeof updatedContext === 'function' + ? updatedContext(OpenFeature.getContext(options?.defaultContext ? undefined : domain)) + : updatedContext; + + if (previousContext.current !== resolvedContext) { if (!domain || options?.defaultContext) { - OpenFeature.setContext(updatedContext); + OpenFeature.setContext(resolvedContext); } else { - OpenFeature.setContext(domain, updatedContext); + OpenFeature.setContext(domain, resolvedContext); } - previousContext.current = updatedContext; + previousContext.current = resolvedContext; } - }, [domain]); + }, [domain, options?.defaultContext]); return { setContext, diff --git a/packages/react/test/provider.spec.tsx b/packages/react/test/provider.spec.tsx index 91277140e..b5f0731b3 100644 --- a/packages/react/test/provider.spec.tsx +++ b/packages/react/test/provider.spec.tsx @@ -162,17 +162,18 @@ describe('OpenFeatureProvider', () => { }); }); describe('useMutateContext', () => { - const MutateButton = () => { + const MutateButton = ({ setter }: { setter?: (prevContext: EvaluationContext) => EvaluationContext }) => { const { setContext } = useContextMutator(); - return ; + return ; }; - const TestComponent = ({ name }: { name: string }) => { + + const TestComponent = ({ name, setter }: { name: string; setter?: (prevContext: EvaluationContext) => EvaluationContext }) => { const flagValue = useStringFlagValue<'hi' | 'bye' | 'aloha'>(SUSPENSE_FLAG_KEY, 'hi'); return (
- +
{`${name} says ${flagValue}`}
); @@ -304,5 +305,37 @@ describe('OpenFeatureProvider', () => { expect(screen.getByText('Will says aloha')).toBeInTheDocument(); }); + + it('should accept a method taking the previous context', async () => { + const DOMAIN = 'mutate-context-with-function'; + OpenFeature.setProvider(DOMAIN, suspendingProvider(), { done: false }); + + const setter = jest.fn((prevContext: EvaluationContext) => ({ ...prevContext, user: 'bob@flags.com' })); + render( + + {FALLBACK}}> + + + , + ); + + await waitFor(() => { + expect(screen.getByText('Will says hi')).toBeInTheDocument(); + }); + + act(() => { + fireEvent.click(screen.getByText('Update Context')); + }); + await waitFor( + () => { + expect(screen.getByText('Will says aloha')).toBeInTheDocument(); + }, + { timeout: DELAY * 4 }, + ); + + expect(setter).toHaveBeenCalledTimes(1); + expect(setter).toHaveBeenCalledWith({ done: false }); + expect(OpenFeature.getContext(DOMAIN)).toEqual({ done: false, user: 'bob@flags.com' }); + }); }); });