fix: await context change in useContextMutator setContext

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

@@ -41,9 +41,9 @@ export function useContextMutator(options: ContextMutationOptions = { defaultCon
 
         if (previousContext.current !== resolvedContext) {
             if (!domain || options?.defaultContext) {
-                OpenFeature.setContext(resolvedContext);
+                await OpenFeature.setContext(resolvedContext);
             } else {
-                OpenFeature.setContext(domain, resolvedContext);
+                await OpenFeature.setContext(domain, resolvedContext);
             }
             previousContext.current = resolvedContext;
         }

packages/react/test/provider.spec.tsx

@@ -1,5 +1,5 @@
 import type { EvaluationContext} from '@openfeature/web-sdk';
-import { InMemoryProvider, OpenFeature } from '@openfeature/web-sdk';
+import { InMemoryProvider, OpenFeature, ProviderEvents } from '@openfeature/web-sdk';
 import '@testing-library/jest-dom'; // see: https://testing-library.com/docs/react-testing-library/setup
 import { render, renderHook, screen, waitFor, fireEvent, act } from '@testing-library/react';
 import * as React from 'react';
@@ -164,8 +164,18 @@ describe('OpenFeatureProvider', () => {
   describe('useMutateContext', () => {
     const MutateButton = ({ setter }: { setter?: (prevContext: EvaluationContext) => EvaluationContext }) => {
       const { setContext } = useContextMutator();
+      const [loading, setLoading] = React.useState(false);
 
-      return <button onClick={() => setContext(setter ?? { user: '[email protected]' })}>Update Context</button>;
+      return (
+        <button
+          onClick={() => {
+            setLoading(true);
+            setContext(setter ?? { user: '[email protected]' }).finally(() => setLoading(false));
+          }}
+        >
+          {loading ? 'Updating context...' : 'Update Context'}
+        </button>
+      );
     };
 
     const TestComponent = ({ name, setter }: { name: string; setter?: (prevContext: EvaluationContext) => EvaluationContext }) => {
@@ -182,6 +192,10 @@ describe('OpenFeatureProvider', () => {
     it('should update context when a domain is set', async () => {
       const DOMAIN = 'mutate-context-tests';
       OpenFeature.setProvider(DOMAIN, suspendingProvider());
+
+      const changed = jest.fn();
+      OpenFeature.getClient(DOMAIN).addHandler(ProviderEvents.ContextChanged, changed);
+
       render(
         <OpenFeatureProvider domain={DOMAIN}>
           <React.Suspense fallback={<div>{FALLBACK}</div>}>
@@ -197,12 +211,17 @@ describe('OpenFeatureProvider', () => {
       act(() => {
         fireEvent.click(screen.getByText('Update Context'));
       });
+      expect(screen.getByText('Updating context...')).toBeInTheDocument();
+
       await waitFor(
         () => {
-          expect(screen.getByText('Will says aloha')).toBeInTheDocument();
+          expect(screen.getByText('Update Context')).toBeInTheDocument();
         },
-        { timeout: DELAY * 4 },
+        { timeout: DELAY * 2 },
       );
+      expect(changed).toHaveBeenCalledTimes(1);
+
+      expect(screen.getByText('Will says aloha')).toBeInTheDocument();
     });
 
     it('should update nested contexts', async () => {
@@ -231,18 +250,17 @@ describe('OpenFeatureProvider', () => {
         // Click the Update context button in Todds domain
         fireEvent.click(screen.getAllByText('Update Context')[1]);
       });
+      expect(screen.getByText('Updating context...')).toBeInTheDocument();
+
       await waitFor(
         () => {
-          expect(screen.getByText('Todd says aloha')).toBeInTheDocument();
-        },
-        { timeout: DELAY * 4 },
-      );
-      await waitFor(
-        () => {
-          expect(screen.getByText('Will says hi')).toBeInTheDocument();
+          expect(screen.getAllByText('Update Context')).toHaveLength(2);
         },
-        { timeout: DELAY * 4 },
+        { timeout: DELAY * 2 },
       );
+
+      expect(screen.getByText('Todd says aloha')).toBeInTheDocument();
+      expect(screen.getByText('Will says hi')).toBeInTheDocument();
     });
 
     it('should update nested global contexts', async () => {
@@ -296,13 +314,16 @@ describe('OpenFeatureProvider', () => {
         // Click the Update context button in Todds domain
         fireEvent.click(screen.getAllByText('Update Context')[1]);
       });
+      expect(screen.getByText('Updating context...')).toBeInTheDocument();
+
       await waitFor(
         () => {
-          expect(screen.getByText('Todd likes to Frown')).toBeInTheDocument();
+          expect(screen.getAllByText('Update Context')).toHaveLength(2);
         },
-        { timeout: DELAY * 4 },
+        { timeout: DELAY * 2 },
       );
 
+      expect(screen.getByText('Todd likes to Frown')).toBeInTheDocument();
       expect(screen.getByText('Will says aloha')).toBeInTheDocument();
     });
 
@@ -326,13 +347,17 @@ describe('OpenFeatureProvider', () => {
       act(() => {
         fireEvent.click(screen.getByText('Update Context'));
       });
+      expect(screen.getByText('Updating context...')).toBeInTheDocument();
+
       await waitFor(
         () => {
-          expect(screen.getByText('Will says aloha')).toBeInTheDocument();
+          expect(screen.getByText('Update Context')).toBeInTheDocument();
         },
-        { timeout: DELAY * 4 },
+        { timeout: DELAY * 2 },
       );
 
+      expect(screen.getByText('Will says aloha')).toBeInTheDocument();
+
       expect(setter).toHaveBeenCalledTimes(1);
       expect(setter).toHaveBeenCalledWith({ done: false });
       expect(OpenFeature.getContext(DOMAIN)).toEqual({ done: false, user: '[email protected]' });