Update MCPConnection.vitest.ts

Author: Patrick-Erichsen

core/context/mcp/MCPConnection.vitest.ts

@@ -2,6 +2,13 @@ import { Client } from "@modelcontextprotocol/sdk/client/index.js";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 import MCPConnection from "./MCPConnection";
 
+// Mock the shell path utility
+vi.mock("../../util/shellPath", () => ({
+  getEnvPathFromUserShell: vi
+    .fn()
+    .mockResolvedValue("/usr/local/bin:/usr/bin:/bin"),
+}));
+
 describe("MCPConnection", () => {
   beforeEach(() => {
     vi.restoreAllMocks();
@@ -76,7 +83,7 @@ describe("MCPConnection", () => {
       expect(conn.status).toBe("not-connected");
     });
 
-    it("should throw on invalid transport type", () => {
+    it("should throw on invalid transport type", async () => {
       const options = {
         name: "test-mcp",
         id: "test-id",
@@ -85,9 +92,14 @@ describe("MCPConnection", () => {
         } as any,
       };
 
-      expect(() => new MCPConnection(options)).toThrow(
-        "Unsupported transport type: invalid",
-      );
+      const conn = new MCPConnection(options);
+      const abortController = new AbortController();
+
+      // The validation now happens during connectClient, not constructor
+      await conn.connectClient(false, abortController.signal);
+
+      expect(conn.status).toBe("error");
+      expect(conn.errors[0]).toContain("Unsupported transport type: invalid");
     });
   });
 
@@ -181,6 +193,15 @@ describe("MCPConnection", () => {
           () => new Promise((resolve) => setTimeout(resolve, 1000)),
         );
 
+      // Mock the required methods for successful connection
+      const mockGetServerCapabilities = vi
+        .spyOn(Client.prototype, "getServerCapabilities")
+        .mockReturnValue({
+          resources: {},
+          tools: {},
+          prompts: {},
+        });
+
       const abortController = new AbortController();
       await conn.connectClient(false, abortController.signal);
 
@@ -189,7 +210,7 @@ describe("MCPConnection", () => {
     });
 
     it("should handle connection timeout", async () => {
-      const conn = new MCPConnection({ ...options, timeout: 1 });
+      const conn = new MCPConnection({ ...options, timeout: 50 });
       const mockConnect = vi
         .spyOn(Client.prototype, "connect")
         .mockImplementation(
@@ -201,7 +222,9 @@ describe("MCPConnection", () => {
 
       expect(conn.status).toBe("error");
       expect(conn.errors[0]).toContain("Failed to connect");
-      expect(mockConnect).toHaveBeenCalled();
+      // The connection should timeout before connect is called due to transport construction
+      // Since transport construction happens first, and we're not mocking that,
+      // the timeout will happen during transport construction or connect attempt
     });
 
     it("should handle already connected state", async () => {