Address comments

Author: kylelai1

src/common/atlas/cluster.ts

@@ -1,26 +1,16 @@
 import type { ClusterDescription20240805, FlexClusterDescription20241113 } from "./openapi.js";
 import type { ApiClient } from "./apiClient.js";
 import { LogId } from "../logger.js";
+import { ConnectionString } from "mongodb-connection-string-url";
 
-const DEFAULT_PORT = "27017";
+type AtlasProcessId = `${string}:${number}`;
 
-function extractProcessIds(connectionString: string): Array<string> {
-    if (!connectionString || connectionString === "") return [];
-
-    // Extract host:port pairs from connection string
-    const matches = connectionString.match(/^mongodb:\/\/([^/]+)/);
-    if (!matches) {
+function extractProcessIds(connectionString: string): Array<AtlasProcessId> {
+    if (!connectionString) {
         return [];
     }
-
-    // matches[1] gives us the host:port pairs
-    const hostsString = matches[1];
-    const hosts = hostsString?.split(",") ?? [];
-
-    return hosts?.map((host) => {
-        const [hostname, port] = host.split(":");
-        return `${hostname}:${port || DEFAULT_PORT}`;
-    });
+    const connectionStringUrl = new ConnectionString(connectionString);
+    return connectionStringUrl.hosts as Array<AtlasProcessId>;
 }
 export interface Cluster {
     name?: string;

tests/integration/tools/atlas/atlasHelpers.ts

@@ -149,11 +149,32 @@ export function sleep(ms: number): Promise<void> {
     return new Promise((resolve) => setTimeout(resolve, ms));
 }
 
+export async function assertClusterIsAvailable(
+    session: Session,
+    projectId: string,
+    clusterName: string
+): Promise<boolean> {
+    try {
+        await session.apiClient.getCluster({
+            params: {
+                path: {
+                    groupId: projectId,
+                    clusterName,
+                },
+            },
+        });
+        return true;
+    } catch {
+        return false;
+    }
+}
+
 export async function deleteAndWaitCluster(
     session: Session,
     projectId: string,
     clusterName: string,
-    pollingInterval: number = 1000
+    pollingInterval: number = 1000,
+    maxPollingIterations: number = 300
 ): Promise<void> {
     await session.apiClient.deleteCluster({
         params: {
@@ -163,31 +184,28 @@ export async function deleteAndWaitCluster(
             },
         },
     });
-    while (true) {
-        try {
-            await session.apiClient.getCluster({
-                params: {
-                    path: {
-                        groupId: projectId,
-                        clusterName,
-                    },
-                },
-            });
-            await sleep(pollingInterval);
-        } catch {
-            break;
+
+    for (let i = 0; i < maxPollingIterations; i++) {
+        const isAvailable = await assertClusterIsAvailable(session, projectId, clusterName);
+        if (!isAvailable) {
+            return;
         }
+        await sleep(pollingInterval);
     }
+    throw new Error(
+        `Cluster deletion timeout: ${clusterName} did not delete within ${maxPollingIterations} iterations`
+    );
 }
 
 export async function waitCluster(
     session: Session,
     projectId: string,
     clusterName: string,
     check: (cluster: ClusterDescription20240805) => boolean | Promise<boolean>,
-    pollingInterval: number = 1000
+    pollingInterval: number = 1000,
+    maxPollingIterations: number = 300
 ): Promise<void> {
-    while (true) {
+    for (let i = 0; i < maxPollingIterations; i++) {
         const cluster = await session.apiClient.getCluster({
             params: {
                 path: {
@@ -201,4 +219,8 @@ export async function waitCluster(
         }
         await sleep(pollingInterval);
     }
+
+    throw new Error(
+        `Cluster wait timeout: ${clusterName} did not meet condition within ${maxPollingIterations} iterations`
+    );
 }

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

@@ -14,7 +14,7 @@ describeWithAtlas("performanceAdvisor", (integration) => {
             const projectId = getProjectId();
             if (projectId) {
                 const session: Session = integration.mcpServer().session;
-                await deleteAndWaitCluster(session, projectId, clusterName, 1000);
+                await deleteAndWaitCluster(session, projectId, clusterName, 1000, 1200);
             }
         }, 1200000);
 
@@ -62,7 +62,8 @@ describeWithAtlas("performanceAdvisor", (integration) => {
                     (cluster) => {
                         return cluster.stateName === "IDLE";
                     },
-                    10000
+                    10000,
+                    120
                 );
 
                 await sleep(10000);