diff --git a/features/update/async_accepted/feature.ts b/features/update/async_accepted/feature.ts
index bde989b9..54534278 100644
--- a/features/update/async_accepted/feature.ts
+++ b/features/update/async_accepted/feature.ts
@@ -3,6 +3,7 @@ import { Feature } from '@temporalio/harness';
 import * as assert from 'assert';
 import ms from 'ms';
 import { Duration, StringValue } from '@temporalio/common';
+import { WorkflowUpdateStage } from '@temporalio/client';
 
 const myUpdate = wf.defineUpdate<void, [Duration, boolean]>('myUpdate');
 const requestedSleep = '2s';
@@ -10,7 +11,7 @@ const requestedSleep = '2s';
 export const feature = new Feature({
   workflow,
   checkResult: async (_, handle) => {
-    const timeToAccept = await time(handle.startUpdate(myUpdate, { args: [requestedSleep, false] }));
+    const timeToAccept = await time(handle.startUpdate(myUpdate, { args: [requestedSleep, false], waitForStage: WorkflowUpdateStage.ACCEPTED }));
     const timeToComplete = await time(handle.executeUpdate(myUpdate, { args: [requestedSleep, false] }));
     assert.equal(
       ms(timeToAccept) < ms(requestedSleep),
diff --git a/features/update/basic_async/feature.ts b/features/update/basic_async/feature.ts
index 6f23769e..db8c85bb 100644
--- a/features/update/basic_async/feature.ts
+++ b/features/update/basic_async/feature.ts
@@ -1,4 +1,4 @@
-import { WorkflowUpdateFailedError } from '@temporalio/client';
+import { WorkflowUpdateFailedError, WorkflowUpdateStage } from '@temporalio/client';
 import { Feature } from '@temporalio/harness';
 import * as wf from '@temporalio/workflow';
 import * as assert from 'assert';
@@ -29,7 +29,7 @@ export async function workflow(): Promise<string> {
 export const feature = new Feature({
   workflow,
   checkResult: async (runner, handle) => {
-    const badUpdateHandle = await handle.startUpdate(myUpdate, { args: ['invalid-arg'] });
+    const badUpdateHandle = await handle.startUpdate(myUpdate, { args: ['invalid-arg'], waitForStage: WorkflowUpdateStage.ACCEPTED });
     try {
       await badUpdateHandle.result();
       throw 'Expected update to fail';
@@ -39,7 +39,7 @@ export const feature = new Feature({
       }
     }
 
-    const updateHandle = await handle.startUpdate(myUpdate, { args: ['update-arg'] });
+    const updateHandle = await handle.startUpdate(myUpdate, { args: ['update-arg'], waitForStage: WorkflowUpdateStage.ACCEPTED });
     const updateResult = await updateHandle.result();
     assert.equal(updateResult, 'update-result');
     const workflowResult = await runner.waitForRunResult(handle);