flaky test fix

Author: THardy98

packages/test/src/helpers-integration.ts

@@ -285,6 +285,16 @@ export function configurableHelpers<T>(
   };
 }
 
+/**
+ * Set the pause state of an activity and wait for the state to be reflected in the workflow description.
+ *
+ * This function handles a race condition where an activity might complete (due to catching a cancellation)
+ * before we can verify that it's been paused. In that case, if the activity is no longer in pendingActivities
+ * and we're trying to pause it, we consider the operation successful rather than timing out.
+ *
+ * This is particularly important for tests where the activity can complete immediately upon receiving a pause
+ * cancellation.
+ */
 export async function setActivityPauseState(handle: WorkflowHandle, activityId: string, pause: boolean): Promise<void> {
   const desc = await handle.describe();
   const req = {
@@ -303,7 +313,14 @@ export async function setActivityPauseState(handle: WorkflowHandle, activityId:
   await waitUntil(async () => {
     const { raw } = await handle.describe();
     const activityInfo = raw.pendingActivities?.find((act) => act.activityId === activityId);
-    if (!activityInfo) return false;
+
+    // If we're trying to pause and the activity is no longer in pendingActivities,
+    // it might have completed due to catching the cancellation. Consider this a success.
+    if (!activityInfo) {
+      // Only consider this a success for pause operations, not unpause
+      return pause;
+    }
+
     if (pause) {
       return activityInfo.paused ?? false;
     }

packages/test/src/test-integration-workflows.ts

@@ -1517,10 +1517,22 @@ test('Activity pause returns expected cancellation details', async (t) => {
     const handle = await startWorkflow(heartbeatPauseWorkflow, { args: [testActivityId, true, 1] });
 
     await waitUntil(async () => handle.query(activityStartedQuery, 1), 5000);
+    // Wait for activity to send at least one heartbeat before pausing
+    await waitUntil(async () => {
+      const { raw } = await handle.describe();
+      const activityInfo = raw.pendingActivities?.find((act) => act.activityId === testActivityId);
+      return activityInfo !== undefined;
+    }, 5000);
     await setActivityPauseState(handle, testActivityId, true);
     await handle.signal(proceedSignal);
 
     await waitUntil(async () => handle.query(activityStartedQuery, 2), 5000);
+    // Wait for activity to be present before pausing
+    await waitUntil(async () => {
+      const { raw } = await handle.describe();
+      const activityInfo = raw.pendingActivities?.find((act) => act.activityId === `${testActivityId}-2`);
+      return activityInfo !== undefined;
+    }, 5000);
     await setActivityPauseState(handle, `${testActivityId}-2`, true);
     await handle.signal(proceedSignal);