Fix several IDE mode integration tests (#6408)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

2 files changed, 96 insertions, 14 deletions

diff --git a/integration-tests/ide-client.test.ts b/integration-tests/ide-client.test.ts
index 10f8a7eb..ce1290b7 100644
--- a/integration-tests/ide-client.test.ts
+++ b/integration-tests/ide-client.test.ts
@@ -9,23 +9,34 @@ import * as fs from 'node:fs';
 import * as os from 'node:os';
 import * as path from 'node:path';
 import * as net from 'node:net';
+import * as child_process from 'node:child_process';
 import { IdeClient } from '../packages/core/src/ide/ide-client.js';
-import { getIdeProcessId } from '../packages/core/src/ide/process-utils.js';
-import { spawn, ChildProcess } from 'child_process';
+
+import { TestMcpServer } from './test-mcp-server.js';
 
 describe('IdeClient', () => {
   it('reads port from file and connects', async () => {
-    const port = 12345;
-    const pid = await getIdeProcessId();
+    const server = new TestMcpServer();
+    const port = await server.start();
+    const pid = process.pid;
     const portFile = path.join(os.tmpdir(), `gemini-ide-server-${pid}.json`);
     fs.writeFileSync(portFile, JSON.stringify({ port }));
+    process.env['GEMINI_CLI_IDE_WORKSPACE_PATH'] = process.cwd();
+    process.env['TERM_PROGRAM'] = 'vscode';
 
     const ideClient = IdeClient.getInstance();
     await ideClient.connect();
 
-    expect(ideClient.getConnectionStatus().status).not.toBe('disconnected');
+    expect(ideClient.getConnectionStatus()).toEqual({
+      status: 'connected',
+      details: undefined,
+    });
 
     fs.unlinkSync(portFile);
+    await server.stop();
+    delete process.env['GEMINI_CLI_IDE_WORKSPACE_PATH'];
+    // Reset instance
+    IdeClient.instance = undefined;
   });
 });
 
@@ -44,24 +55,25 @@ const getFreePort = (): Promise<number> => {
 };
 
 describe('IdeClient fallback connection logic', () => {
-  let server: net.Server;
+  let server: TestMcpServer;
   let envPort: number;
   let pid: number;
   let portFile: string;
 
   beforeEach(async () => {
-    pid = await getIdeProcessId();
+    pid = process.pid;
     portFile = path.join(os.tmpdir(), `gemini-ide-server-${pid}.json`);
-    envPort = await getFreePort();
-    server = net.createServer().listen(envPort);
+    server = new TestMcpServer();
+    envPort = await server.start();
     process.env['GEMINI_CLI_IDE_SERVER_PORT'] = String(envPort);
+    process.env['TERM_PROGRAM'] = 'vscode';
     process.env['GEMINI_CLI_IDE_WORKSPACE_PATH'] = process.cwd();
     // Reset instance
     IdeClient.instance = undefined;
   });
 
-  afterEach(() => {
-    server.close();
+  afterEach(async () => {
+    await server.stop();
     delete process.env['GEMINI_CLI_IDE_SERVER_PORT'];
     delete process.env['GEMINI_CLI_IDE_WORKSPACE_PATH'];
     if (fs.existsSync(portFile)) {
@@ -78,7 +90,10 @@ describe('IdeClient fallback connection logic', () => {
     const ideClient = IdeClient.getInstance();
     await ideClient.connect();
 
-    expect(ideClient.getConnectionStatus().status).toBe('connected');
+    expect(ideClient.getConnectionStatus()).toEqual({
+      status: 'connected',
+      details: undefined,
+    });
   });
 
   it('falls back to env var when connection with port from file fails', async () => {
@@ -89,7 +104,10 @@ describe('IdeClient fallback connection logic', () => {
     const ideClient = IdeClient.getInstance();
     await ideClient.connect();
 
-    expect(ideClient.getConnectionStatus().status).toBe('connected');
+    expect(ideClient.getConnectionStatus()).toEqual({
+      status: 'connected',
+      details: undefined,
+    });
   });
 });
 
@@ -107,7 +125,7 @@ describe('getIdeProcessId', () => {
     // so that we can check that getIdeProcessId returns the pid of the parent
     const parentPid = process.pid;
     const output = await new Promise<string>((resolve, reject) => {
-      child = spawn(
+      child = child_process.spawn(
         'node',
         [
           '-e',
diff --git a/integration-tests/test-mcp-server.ts b/integration-tests/test-mcp-server.ts
new file mode 100644
index 00000000..121d6ed0
--- /dev/null
+++ b/integration-tests/test-mcp-server.ts
@@ -0,0 +1,64 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
+import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
+import express from 'express';
+import { type Server as HTTPServer } from 'node:http';
+
+import { randomUUID } from 'node:crypto';
+
+export class TestMcpServer {
+  private server: HTTPServer | undefined;
+
+  async start(): Promise<number> {
+    const app = express();
+    app.use(express.json());
+    const mcpServer = new McpServer(
+      {
+        name: 'test-mcp-server',
+        version: '1.0.0',
+      },
+      { capabilities: {} },
+    );
+
+    const transport = new StreamableHTTPServerTransport({
+      sessionIdGenerator: () => randomUUID(),
+    });
+    mcpServer.connect(transport);
+
+    app.post('/mcp', async (req, res) => {
+      await transport.handleRequest(req, res, req.body);
+    });
+
+    return new Promise((resolve, reject) => {
+      this.server = app.listen(0, () => {
+        const address = this.server!.address();
+        if (address && typeof address !== 'string') {
+          resolve(address.port);
+        } else {
+          reject(new Error('Could not determine server port.'));
+        }
+      });
+      this.server.on('error', reject);
+    });
+  }
+
+  async stop(): Promise<void> {
+    if (this.server) {
+      await new Promise<void>((resolve, reject) => {
+        this.server!.close((err?: Error) => {
+          if (err) {
+            reject(err);
+          } else {
+            resolve();
+          }
+        });
+      });
+      this.server = undefined;
+    }
+  }
+}