diff --git a/server/src/__tests__/heartbeat-retry-scheduling.test.ts b/server/src/__tests__/heartbeat-retry-scheduling.test.ts
index 3dada976b57..b6b197f1485 100644
--- a/server/src/__tests__/heartbeat-retry-scheduling.test.ts
+++ b/server/src/__tests__/heartbeat-retry-scheduling.test.ts
@@ -219,6 +219,7 @@ describeEmbeddedPostgres("heartbeat bounded retry scheduling", () => {
     expect(failedRun?.status).toBe("failed");
     expect(failedRun?.errorCode).toBe("provider_quota");
     expect((failedRun?.resultJson as Record<string, unknown> | null)?.errorFamily).toBe("provider_quota");
+    await heartbeat.waitForRunExecutionDrain(run!.id);
 
     await expect
       .poll(