Allow method to be passed into setContext hook

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

Author: MattIPv4

packages/react/src/context/use-context-mutator.ts

@@ -21,7 +21,7 @@ export type ContextMutation = {
    * @param updatedContext
    * @returns Promise for awaiting the context update
    */
-  setContext: (updatedContext: EvaluationContext) => Promise<void>;
+  setContext: (updatedContext: EvaluationContext | ((previousContext: EvaluationContext) => EvaluationContext)) => Promise<void>;
 };
 
 /**
@@ -34,14 +34,18 @@ 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 | ((previousContext: 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]);