feat: allow method to be passed into setContext hook (#1301)

## This PR

Allows for `setContext` returned by the `useContextMutator` hook to be
passed a method that returns the new context, with the current
(previous) context passed into it. This aligns the signature of
`setContext` to the setter of a regular `useState` hook, making it
easier to conditionally update the context without needing to keep track
of the domain separately from the React context.

---------

Signed-off-by: MattIPv4 <[email protected]>

Author: MattIPv4

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<void>;
+  setContext: (updatedContext: EvaluationContext | ((currentContext: EvaluationContext) => EvaluationContext)) => Promise<void>;
 };
 
 /**
@@ -34,16 +34,20 @@ export function useContextMutator(options: ContextMutationOptions = { defaultCon
     const { domain } = useContext(Context) || {};
     const previousContext = useRef<null | EvaluationContext>(null);
 
-    const setContext = useCallback(async (updatedContext: EvaluationContext) => {
-        if (previousContext.current !== updatedContext) {
+    const setContext = useCallback(async (updatedContext: EvaluationContext | ((currentContext: EvaluationContext) => EvaluationContext)): Promise<void> => {
+        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,

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 <button onClick={() => setContext({ user: '[email protected]' })}>Update Context</button>;
+      return <button onClick={() => setContext(setter ?? { user: '[email protected]' })}>Update Context</button>;
     };
-    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 (
         <div>
-          <MutateButton />
+          <MutateButton setter={setter} />
           <div>{`${name} says ${flagValue}`}</div>
         </div>
       );
@@ -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: '[email protected]' }));
+      render(
+        <OpenFeatureProvider domain={DOMAIN}>
+          <React.Suspense fallback={<div>{FALLBACK}</div>}>
+            <TestComponent name="Will" setter={setter} />
+          </React.Suspense>
+        </OpenFeatureProvider>,
+      );
+
+      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: '[email protected]' });
+    });
   });
 });