chore: add retries to setup in case the environment is unstable

Author: kmruiz

tests/integration/tools/atlas/accessLists.test.ts

@@ -1,6 +1,6 @@
-import { describeWithAtlas, withProject } from "./atlasHelpers.js";
+import { afterAllWithRetry, beforeAllWithRetry, describeWithAtlas, withProject } from "./atlasHelpers.js";
 import { expectDefined, getResponseElements } from "../../helpers.js";
-import { afterAll, beforeAll, describe, expect, it } from "vitest";
+import { describe, expect, it } from "vitest";
 import { ensureCurrentIpInAccessList } from "../../../../src/common/atlas/accessListUtils.js";
 
 function generateRandomIp(): string {
@@ -17,13 +17,13 @@ describeWithAtlas("ip access lists", (integration) => {
         const cidrBlocks = [generateRandomIp() + "/16", generateRandomIp() + "/24"];
         const values = [...ips, ...cidrBlocks];
 
-        beforeAll(async () => {
+        beforeAllWithRetry(async () => {
             const apiClient = integration.mcpServer().session.apiClient;
             const ipInfo = await apiClient.getIpInfo();
             values.push(ipInfo.currentIpv4Address);
         });
 
-        afterAll(async () => {
+        afterAllWithRetry(async () => {
             const apiClient = integration.mcpServer().session.apiClient;
 
             const projectId = getProjectId();

tests/integration/tools/atlas/atlasHelpers.ts

@@ -4,7 +4,7 @@ import type { ApiClient } from "../../../../src/common/atlas/apiClient.js";
 import type { IntegrationTest } from "../../helpers.js";
 import { setupIntegrationTest, defaultTestConfig, defaultDriverOptions } from "../../helpers.js";
 import type { SuiteCollector } from "vitest";
-import { afterAll, beforeAll, describe } from "vitest";
+import { afterAll, beforeEach, describe } from "vitest";
 
 export type IntegrationTestFunction = (integration: IntegrationTest) => void;
 
@@ -36,19 +36,13 @@ export function withProject(integration: IntegrationTest, fn: ProjectTestFunctio
     return describe("with project", () => {
         let projectId: string = "";
 
-        beforeAll(async () => {
+        beforeAllWithRetry(async () => {
             const apiClient = integration.mcpServer().session.apiClient;
-
-            try {
-                const group = await createProject(apiClient);
-                projectId = group.id;
-            } catch (error) {
-                console.error("Failed to create project:", error);
-                throw error;
-            }
+            const group = await createProject(apiClient);
+            projectId = group.id;
         });
 
-        afterAll(async () => {
+        afterAllWithRetry(async () => {
             const apiClient = integration.mcpServer().session.apiClient;
             if (projectId) {
                 // projectId may be empty if beforeAll failed.
@@ -70,6 +64,63 @@ export function withProject(integration: IntegrationTest, fn: ProjectTestFunctio
     });
 }
 
+export function beforeAllWithRetry(fixture: () => Promise<void>): void {
+    beforeEach(async () => {
+        const MAX_SETUP_ATTEMPTS = 10;
+        const SETUP_BACKOFF_MS = 10;
+        let lastError: Error | undefined = undefined;
+
+        for (let attempt = 0; attempt < MAX_SETUP_ATTEMPTS; attempt++) {
+            try {
+                await fixture();
+                lastError = undefined;
+                break;
+            } catch (error: unknown) {
+                if (error instanceof Error) {
+                    lastError = error;
+                } else {
+                    lastError = new Error(String(error));
+                }
+
+                console.error("beforeAll(attempt:", attempt, "):", error);
+                await new Promise((resolve) => setTimeout(resolve, SETUP_BACKOFF_MS * attempt));
+            }
+        }
+
+        if (lastError) {
+            throw lastError;
+        }
+    });
+}
+
+export function afterAllWithRetry(fixture: () => Promise<void>): void {
+    afterAll(async () => {
+        const MAX_SETUP_ATTEMPTS = 10;
+        const SETUP_BACKOFF_MS = 10;
+        let lastError: Error | undefined = undefined;
+
+        for (let attempt = 0; attempt < MAX_SETUP_ATTEMPTS; attempt++) {
+            try {
+                await fixture();
+                lastError = undefined;
+                break;
+            } catch (error) {
+                if (error instanceof Error) {
+                    lastError = error;
+                } else {
+                    lastError = new Error(String(error));
+                }
+                console.error("afterAll(attempt:", attempt, "):", error);
+                await new Promise((resolve) => setTimeout(resolve, SETUP_BACKOFF_MS * attempt));
+            }
+        }
+
+        if (lastError) {
+            throw lastError;
+        }
+    });
+}
+
 export function parseTable(text: string): Record<string, string>[] {
     const data = text
         .split("\n")

tests/integration/tools/atlas/clusters.test.ts

@@ -1,8 +1,8 @@
 import type { Session } from "../../../../src/common/session.js";
 import { expectDefined, getResponseElements } from "../../helpers.js";
-import { describeWithAtlas, withProject, randomId } from "./atlasHelpers.js";
+import { describeWithAtlas, withProject, randomId, afterAllWithRetry, beforeAllWithRetry } from "./atlasHelpers.js";
 import type { ClusterDescription20240805 } from "../../../../src/common/atlas/openapi.js";
-import { afterAll, beforeAll, describe, expect, it } from "vitest";
+import { describe, expect, it } from "vitest";
 
 function sleep(ms: number): Promise<void> {
     return new Promise((resolve) => setTimeout(resolve, ms));
@@ -60,7 +60,7 @@ describeWithAtlas("clusters", (integration) => {
     withProject(integration, ({ getProjectId }) => {
         const clusterName = "ClusterTest-" + randomId;
 
-        afterAll(async () => {
+        afterAllWithRetry(async () => {
             const projectId = getProjectId();
             if (projectId) {
                 const session: Session = integration.mcpServer().session;
@@ -160,7 +160,7 @@ describeWithAtlas("clusters", (integration) => {
         });
 
         describe("atlas-connect-cluster", () => {
-            beforeAll(async () => {
+            beforeAllWithRetry(async () => {
                 const projectId = getProjectId();
                 await waitCluster(integration.mcpServer().session, projectId, clusterName, (cluster) => {
                     return (

tests/integration/tools/atlas/projects.test.ts

@@ -1,14 +1,14 @@
 import { ObjectId } from "mongodb";
-import { parseTable, describeWithAtlas } from "./atlasHelpers.js";
+import { parseTable, describeWithAtlas, afterAllWithRetry } from "./atlasHelpers.js";
 import { expectDefined, getDataFromUntrustedContent, getResponseElements } from "../../helpers.js";
-import { afterAll, describe, expect, it } from "vitest";
+import { describe, expect, it } from "vitest";
 
 const randomId = new ObjectId().toString();
 
 describeWithAtlas("projects", (integration) => {
     const projName = "testProj-" + randomId;
 
-    afterAll(async () => {
+    afterAllWithRetry(async () => {
         const session = integration.mcpServer().session;
 
         const projects = await session.apiClient.listProjects();