Clean up processes in integration tests (#15102)

integration-tests/ctrl-c-exit.test.ts

@@ -4,13 +4,20 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { describe, it, expect } from 'vitest';
+import { describe, it, expect, beforeEach, afterEach } from 'vitest';
 import * as os from 'node:os';
 import { TestRig } from './test-helper.js';
 
 describe('Ctrl+C exit', () => {
+  let rig: TestRig;
+
+  beforeEach(() => {
+    rig = new TestRig();
+  });
+
+  afterEach(async () => await rig.cleanup());
+
   it('should exit gracefully on second Ctrl+C', async () => {
-    const rig = new TestRig();
     await rig.setup('should exit gracefully on second Ctrl+C', {
       settings: { tools: { useRipgrep: false } },
     });

integration-tests/extensions-install.test.ts

@@ -4,7 +4,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { describe, expect, it } from 'vitest';
+import { describe, expect, it, beforeEach, afterEach } from 'vitest';
 import { TestRig } from './test-helper.js';
 import { writeFileSync } from 'node:fs';
 import { join } from 'node:path';
@@ -20,8 +20,15 @@ const extensionUpdate = `{
 }`;
 
 describe('extension install', () => {
+  let rig: TestRig;
+
+  beforeEach(() => {
+    rig = new TestRig();
+  });
+
+  afterEach(async () => await rig.cleanup());
+
   it('installs a local extension, verifies a command, and updates it', async () => {
-    const rig = new TestRig();
     rig.setup('extension install test');
     const testServerPath = join(rig.testDir!, 'gemini-extension.json');
     writeFileSync(testServerPath, extension);
@@ -47,7 +54,6 @@ describe('extension install', () => {
         'uninstall',
         'test-extension-install',
       ]);
-      await rig.cleanup();
     }
   });
 });

integration-tests/extensions-reload.test.ts

@@ -4,7 +4,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { expect, it, describe } from 'vitest';
+import { expect, it, describe, beforeEach, afterEach } from 'vitest';
 import { TestRig } from './test-helper.js';
 import { TestMcpServer } from './test-mcp-server.js';
 import { writeFileSync } from 'node:fs';
@@ -18,6 +18,14 @@ import stripAnsi from 'strip-ansi';
 const itIf = (condition: boolean) => (condition ? it : it.skip);
 
 describe('extension reloading', () => {
+  let rig: TestRig;
+
+  beforeEach(() => {
+    rig = new TestRig();
+  });
+
+  afterEach(async () => await rig.cleanup());
+
   const sandboxEnv = env['GEMINI_SANDBOX'];
   // Fails in linux non-sandbox e2e tests
   // TODO(#14527): Re-enable this once fixed
@@ -43,7 +51,6 @@ describe('extension reloading', () => {
         },
       };
 
-      const rig = new TestRig();
       rig.setup('extension reload test', {
         settings: {
           experimental: { extensionReloading: true },
@@ -145,7 +152,6 @@ describe('extension reloading', () => {
       await serverA.stop();
       await serverB.stop();
       await rig.runCommand(['extensions', 'uninstall', 'test-extension']);
-      await rig.cleanup();
     },
   );
 });

integration-tests/file-system.test.ts

@@ -4,14 +4,21 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { describe, it, expect } from 'vitest';
+import { describe, it, expect, beforeEach, afterEach } from 'vitest';
 import { existsSync } from 'node:fs';
 import * as path from 'node:path';
 import { TestRig, printDebugInfo, validateModelOutput } from './test-helper.js';
 
 describe('file-system', () => {
+  let rig: TestRig;
+
+  beforeEach(() => {
+    rig = new TestRig();
+  });
+
+  afterEach(async () => await rig.cleanup());
+
   it('should be able to read a file', async () => {
-    const rig = new TestRig();
     await rig.setup('should be able to read a file', {
       settings: { tools: { core: ['read_file'] } },
     });
@@ -41,7 +48,6 @@ describe('file-system', () => {
   });
 
   it('should be able to write a file', async () => {
-    const rig = new TestRig();
     await rig.setup('should be able to write a file', {
       settings: { tools: { core: ['write_file', 'replace', 'read_file'] } },
     });
@@ -98,7 +104,6 @@ describe('file-system', () => {
   });
 
   it('should correctly handle file paths with spaces', async () => {
-    const rig = new TestRig();
     await rig.setup('should correctly handle file paths with spaces', {
       settings: { tools: { core: ['write_file', 'read_file'] } },
     });
@@ -122,7 +127,6 @@ describe('file-system', () => {
   });
 
   it('should perform a read-then-write sequence', async () => {
-    const rig = new TestRig();
     await rig.setup('should perform a read-then-write sequence', {
       settings: { tools: { core: ['read_file', 'replace', 'write_file'] } },
     });
@@ -159,7 +163,6 @@ describe('file-system', () => {
   });
 
   it.skip('should replace multiple instances of a string', async () => {
-    const rig = new TestRig();
     rig.setup('should replace multiple instances of a string');
     const fileName = 'ambiguous.txt';
     const fileContent = 'Hey there, \ntest line\ntest line';
@@ -211,7 +214,6 @@ describe('file-system', () => {
   });
 
   it('should fail safely when trying to edit a non-existent file', async () => {
-    const rig = new TestRig();
     await rig.setup(
       'should fail safely when trying to edit a non-existent file',
       { settings: { tools: { core: ['read_file', 'replace'] } } },

integration-tests/flicker.test.ts

@@ -4,36 +4,39 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { describe, it, expect } from 'vitest';
+import { describe, it, expect, beforeEach, afterEach } from 'vitest';
 import { TestRig } from './test-helper.js';
 import { join } from 'node:path';
 
 describe('Flicker Detector', () => {
+  let rig: TestRig;
+
+  beforeEach(() => {
+    rig = new TestRig();
+  });
+
+  afterEach(async () => await rig.cleanup());
+
   it('should not detect a flicker under the max height budget', async () => {
-    const rig = new TestRig();
     rig.setup('flicker-detector-test', {
       fakeResponsesPath: join(
         import.meta.dirname,
         'flicker-detector.max-height.responses',
       ),
     });
-    try {
+    const run = await rig.runInteractive();
-      const run = await rig.runInteractive();
+    const prompt = 'Tell me a fun fact.';
-      const prompt = 'Tell me a fun fact.';
+    await run.type(prompt);
-      await run.type(prompt);
+    await run.type('\r');
-      await run.type('\r');
 
-      const hasUserPromptEvent = await rig.waitForTelemetryEvent('user_prompt');
+    const hasUserPromptEvent = await rig.waitForTelemetryEvent('user_prompt');
-      expect(hasUserPromptEvent).toBe(true);
+    expect(hasUserPromptEvent).toBe(true);
 
-      const hasSessionCountMetric = await rig.waitForMetric('session.count');
+    const hasSessionCountMetric = await rig.waitForMetric('session.count');
-      expect(hasSessionCountMetric).toBe(true);
+    expect(hasSessionCountMetric).toBe(true);
 
-      // We expect NO flicker event to be found.
+    // We expect NO flicker event to be found.
-      const flickerMetric = rig.readMetric('ui.flicker.count');
+    const flickerMetric = rig.readMetric('ui.flicker.count');
-      expect(flickerMetric).toBeNull();
+    expect(flickerMetric).toBeNull();
-    } finally {
-      await rig.cleanup();
-    }
   });
 });

integration-tests/google_web_search.test.ts

@@ -5,12 +5,19 @@
  */
 
 import { WEB_SEARCH_TOOL_NAME } from '../packages/core/src/tools/tool-names.js';
-import { describe, it, expect } from 'vitest';
+import { describe, it, expect, beforeEach, afterEach } from 'vitest';
 import { TestRig, printDebugInfo, validateModelOutput } from './test-helper.js';
 
-describe(WEB_SEARCH_TOOL_NAME, () => {
+describe('web search tool', () => {
+  let rig: TestRig;
+
+  beforeEach(() => {
+    rig = new TestRig();
+  });
+
+  afterEach(async () => await rig.cleanup());
+
   it('should be able to search the web', async () => {
-    const rig = new TestRig();
     await rig.setup('should be able to search the web', {
       settings: { tools: { core: [WEB_SEARCH_TOOL_NAME] } },
     });

integration-tests/list_directory.test.ts

@@ -4,7 +4,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { describe, it } from 'vitest';
+import { describe, it, beforeEach, afterEach } from 'vitest';
 import {
   TestRig,
   poll,
@@ -15,8 +15,15 @@ import { existsSync } from 'node:fs';
 import { join } from 'node:path';
 
 describe('list_directory', () => {
+  let rig: TestRig;
+
+  beforeEach(() => {
+    rig = new TestRig();
+  });
+
+  afterEach(async () => await rig.cleanup());
+
   it('should be able to list a directory', async () => {
-    const rig = new TestRig();
     await rig.setup('should be able to list a directory', {
       settings: { tools: { core: ['list_directory'] } },
     });

integration-tests/mcp_server_cyclic_schema.test.ts

@@ -23,7 +23,7 @@
 
 import { writeFileSync } from 'node:fs';
 import { join } from 'node:path';
-import { beforeAll, describe, it } from 'vitest';
+import { describe, it, afterEach, beforeEach } from 'vitest';
 import { TestRig } from './test-helper.js';
 
 // Create a minimal MCP server that doesn't require external dependencies
@@ -166,9 +166,15 @@ rpc.send({
 `;
 
 describe('mcp server with cyclic tool schema is detected', () => {
-  const rig = new TestRig();
+  let rig: TestRig;
 
-  beforeAll(async () => {
+  beforeEach(() => {
+    rig = new TestRig();
+  });
+
+  afterEach(async () => await rig.cleanup());
+
+  it('mcp tool list should include tool with cyclic tool schema', async () => {
     // Setup test directory with MCP server configuration
     await rig.setup('cyclic-schema-mcp-server', {
       settings: {
@@ -190,9 +196,7 @@ describe('mcp server with cyclic tool schema is detected', () => {
       const { chmodSync } = await import('node:fs');
       chmodSync(testServerPath, 0o755);
     }
-  });
 
-  it('mcp tool list should include tool with cyclic tool schema', async () => {
     const run = await rig.runInteractive();
 
     await run.type('/mcp list');

integration-tests/mixed-input-crash.test.ts

@@ -4,12 +4,19 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { describe, it, expect } from 'vitest';
+import { describe, it, expect, beforeEach, afterEach } from 'vitest';
 import { TestRig } from './test-helper.js';
 
 describe('mixed input crash prevention', () => {
+  let rig: TestRig;
+
+  beforeEach(() => {
+    rig = new TestRig();
+  });
+
+  afterEach(async () => await rig.cleanup());
+
   it('should not crash when using mixed prompt inputs', async () => {
-    const rig = new TestRig();
     rig.setup('should not crash when using mixed prompt inputs');
 
     // Test: echo "say '1'." | gemini --prompt-interactive="say '2'." say '3'.
@@ -40,7 +47,6 @@ describe('mixed input crash prevention', () => {
   });
 
   it('should provide clear error message for mixed input', async () => {
-    const rig = new TestRig();
     rig.setup('should provide clear error message for mixed input');
 
     try {

integration-tests/read_many_files.test.ts

@@ -4,12 +4,19 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { describe, it, expect } from 'vitest';
+import { describe, it, expect, beforeEach, afterEach } from 'vitest';
 import { TestRig, printDebugInfo, validateModelOutput } from './test-helper.js';
 
 describe('read_many_files', () => {
+  let rig: TestRig;
+
+  beforeEach(() => {
+    rig = new TestRig();
+  });
+
+  afterEach(async () => await rig.cleanup());
+
   it.skip('should be able to read multiple files', async () => {
-    const rig = new TestRig();
     await rig.setup('should be able to read multiple files', {
       settings: { tools: { core: ['read_many_files', 'read_file'] } },
     });
@@ -45,6 +52,5 @@ describe('read_many_files', () => {
 
     // Validate model output - will throw if no output
     validateModelOutput(result, null, 'Read many files test');
-    await rig.cleanup();
   });
 });

integration-tests/replace.test.ts

@@ -4,12 +4,18 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { describe, it, expect } from 'vitest';
+import { describe, it, expect, beforeEach, afterEach } from 'vitest';
 import { TestRig } from './test-helper.js';
 
 describe('replace', () => {
+  let rig: TestRig;
+
+  beforeEach(() => {
+    rig = new TestRig();
+  });
+
+  afterEach(async () => await rig.cleanup());
   it('should be able to replace content in a file', async () => {
-    const rig = new TestRig();
     await rig.setup('should be able to replace content in a file', {
       settings: { tools: { core: ['replace', 'read_file'] } },
     });
@@ -29,7 +35,6 @@ describe('replace', () => {
   });
 
   it.skip('should handle $ literally when replacing text ending with $', async () => {
-    const rig = new TestRig();
     await rig.setup(
       'should handle $ literally when replacing text ending with $',
       { settings: { tools: { core: ['replace', 'read_file'] } } },
@@ -52,7 +57,6 @@ describe('replace', () => {
   });
 
   it.skip('should insert a multi-line block of text', async () => {
-    const rig = new TestRig();
     await rig.setup('should insert a multi-line block of text', {
       settings: { tools: { core: ['replace', 'read_file'] } },
     });
@@ -73,7 +77,6 @@ describe('replace', () => {
   });
 
   it.skip('should delete a block of text', async () => {
-    const rig = new TestRig();
     await rig.setup('should delete a block of text', {
       settings: { tools: { core: ['replace', 'read_file'] } },
     });

integration-tests/run_shell_command.test.ts

@@ -4,7 +4,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { describe, it, expect } from 'vitest';
+import { describe, it, expect, beforeEach, afterEach } from 'vitest';
 import { TestRig, printDebugInfo, validateModelOutput } from './test-helper.js';
 import { getShellConfiguration } from '../packages/core/src/utils/shell-utils.js';
 
@@ -84,8 +84,14 @@ function getChainedEchoCommand(): { allowPattern: string; command: string } {
 }
 
 describe('run_shell_command', () => {
+  let rig: TestRig;
+
+  beforeEach(() => {
+    rig = new TestRig();
+  });
+
+  afterEach(async () => await rig.cleanup());
   it('should be able to run a shell command', async () => {
-    const rig = new TestRig();
     await rig.setup('should be able to run a shell command', {
       settings: { tools: { core: ['run_shell_command'] } },
     });
@@ -119,7 +125,6 @@ describe('run_shell_command', () => {
   });
 
   it('should be able to run a shell command via stdin', async () => {
-    const rig = new TestRig();
     await rig.setup('should be able to run a shell command via stdin', {
       settings: { tools: { core: ['run_shell_command'] } },
     });
@@ -149,7 +154,6 @@ describe('run_shell_command', () => {
   });
 
   it.skip('should run allowed sub-command in non-interactive mode', async () => {
-    const rig = new TestRig();
     await rig.setup('should run allowed sub-command in non-interactive mode');
 
     const testFile = rig.createFile('test.txt', 'Lorem\nIpsum\nDolor\n');
@@ -196,7 +200,6 @@ describe('run_shell_command', () => {
   });
 
   it.skip('should succeed with no parens in non-interactive mode', async () => {
-    const rig = new TestRig();
     await rig.setup('should succeed with no parens in non-interactive mode');
 
     const testFile = rig.createFile('test.txt', 'Lorem\nIpsum\nDolor\n');
@@ -233,7 +236,6 @@ describe('run_shell_command', () => {
   });
 
   it('should succeed with --yolo mode', async () => {
-    const rig = new TestRig();
     await rig.setup('should succeed with --yolo mode', {
       settings: { tools: { core: ['run_shell_command'] } },
     });
@@ -269,7 +271,6 @@ describe('run_shell_command', () => {
   });
 
   it.skip('should work with ShellTool alias', async () => {
-    const rig = new TestRig();
     await rig.setup('should work with ShellTool alias');
 
     const testFile = rig.createFile('test.txt', 'Lorem\nIpsum\nDolor\n');
@@ -317,7 +318,6 @@ describe('run_shell_command', () => {
   // TODO(#11062): Un-skip this once we can make it reliable by using hard coded
   // model responses.
   it.skip('should combine multiple --allowed-tools flags', async () => {
-    const rig = new TestRig();
     await rig.setup('should combine multiple --allowed-tools flags');
 
     const { tool, command } = getLineCountCommand();
@@ -367,7 +367,6 @@ describe('run_shell_command', () => {
   });
 
   it('should reject commands not on the allowlist', async () => {
-    const rig = new TestRig();
     await rig.setup('should reject commands not on the allowlist', {
       settings: { tools: { core: ['run_shell_command'] } },
     });
@@ -437,7 +436,6 @@ describe('run_shell_command', () => {
 
   // TODO(#11966): Deflake this test and re-enable once the underlying race is resolved.
   it.skip('should reject chained commands when only the first segment is allowlisted in non-interactive mode', async () => {
-    const rig = new TestRig();
     await rig.setup(
       'should reject chained commands when only the first segment is allowlisted',
     );
@@ -466,7 +464,6 @@ describe('run_shell_command', () => {
   });
 
   it('should allow all with "ShellTool" and other specific tools', async () => {
-    const rig = new TestRig();
     await rig.setup(
       'should allow all with "ShellTool" and other specific tools',
       {
@@ -516,7 +513,6 @@ describe('run_shell_command', () => {
   });
 
   it('should propagate environment variables to the child process', async () => {
-    const rig = new TestRig();
     await rig.setup('should propagate environment variables', {
       settings: { tools: { core: ['run_shell_command'] } },
     });
@@ -550,7 +546,6 @@ describe('run_shell_command', () => {
   });
 
   it.skip('should run a platform-specific file listing command', async () => {
-    const rig = new TestRig();
     await rig.setup('should run platform-specific file listing');
     const fileName = `test-file-${Math.random().toString(36).substring(7)}.txt`;
     rig.createFile(fileName, 'test content');
@@ -578,7 +573,6 @@ describe('run_shell_command', () => {
   });
 
   it('rejects invalid shell expressions', async () => {
-    const rig = new TestRig();
     await rig.setup('rejects invalid shell expressions', {
       settings: { tools: { core: ['run_shell_command'] } },
     });

integration-tests/save_memory.test.ts

@@ -4,12 +4,19 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { describe, it, expect } from 'vitest';
+import { describe, it, expect, beforeEach, afterEach } from 'vitest';
 import { TestRig, printDebugInfo, validateModelOutput } from './test-helper.js';
 
 describe('save_memory', () => {
+  let rig: TestRig;
+
+  beforeEach(() => {
+    rig = new TestRig();
+  });
+
+  afterEach(async () => await rig.cleanup());
+
   it('should be able to save to memory', async () => {
-    const rig = new TestRig();
     await rig.setup('should be able to save to memory', {
       settings: { tools: { core: ['save_memory'] } },
     });

integration-tests/simple-mcp-server.test.ts

@@ -10,7 +10,7 @@
  * external dependencies, making it compatible with Docker sandbox mode.
  */
 
-import { describe, it, beforeAll, expect } from 'vitest';
+import { describe, it, expect, beforeEach, afterEach } from 'vitest';
 import { TestRig, poll, validateModelOutput } from './test-helper.js';
 import { join } from 'node:path';
 import { writeFileSync } from 'node:fs';
@@ -165,9 +165,15 @@ rpc.send({
 `;
 
 describe('simple-mcp-server', () => {
-  const rig = new TestRig();
+  let rig: TestRig;
 
-  beforeAll(async () => {
+  beforeEach(() => {
+    rig = new TestRig();
+  });
+
+  afterEach(async () => await rig.cleanup());
+
+  it('should add two numbers', async () => {
     // Setup test directory with MCP server configuration
     await rig.setup('simple-mcp-server', {
       settings: {
@@ -209,9 +215,7 @@ describe('simple-mcp-server', () => {
     if (!isReady) {
       throw new Error('MCP server script was not ready in time.');
     }
-  });
 
-  it('should add two numbers', async () => {
     // Test directory is already set up in before hook
     // Just run the command - MCP server config is in settings.json
     const output = await rig.run(

integration-tests/stdin-context.test.ts

@@ -4,12 +4,19 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { describe, it, expect } from 'vitest';
+import { describe, it, expect, beforeEach, afterEach } from 'vitest';
 import { TestRig, printDebugInfo, validateModelOutput } from './test-helper.js';
 
 describe.skip('stdin context', () => {
+  let rig: TestRig;
+
+  beforeEach(() => {
+    rig = new TestRig();
+  });
+
+  afterEach(async () => await rig.cleanup());
+
   it('should be able to use stdin as context for a prompt', async () => {
-    const rig = new TestRig();
     await rig.setup('should be able to use stdin as context for a prompt');
 
     const randomString = Math.random().toString(36).substring(7);
@@ -75,7 +82,6 @@ describe.skip('stdin context', () => {
       even though gemini is intended to run interactively.
     */
 
-    const rig = new TestRig();
     await rig.setup('should exit quickly if stdin stream does not end');
 
     try {

integration-tests/telemetry.test.ts

@@ -4,12 +4,19 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { describe, it, expect } from 'vitest';
+import { describe, it, expect, beforeEach, afterEach } from 'vitest';
 import { TestRig } from './test-helper.js';
 
 describe('telemetry', () => {
+  let rig: TestRig;
+
+  beforeEach(() => {
+    rig = new TestRig();
+  });
+
+  afterEach(async () => await rig.cleanup());
+
   it('should emit a metric and a log event', async () => {
-    const rig = new TestRig();
     rig.setup('should emit a metric and a log event');
 
     // Run a simple command that should trigger telemetry

integration-tests/test-helper.ts

@@ -274,6 +274,7 @@ export class TestRig {
   fakeResponsesPath?: string;
   // Original fake responses file path for rewriting goldens in record mode.
   originalFakeResponsesPath?: string;
+  private _interactiveRuns: InteractiveRun[] = [];
 
   constructor() {
     this.bundlePath = join(__dirname, '..', 'bundle/gemini.js');
@@ -586,6 +587,18 @@ export class TestRig {
   }
 
   async cleanup() {
+    // Kill any interactive runs that are still active
+    for (const run of this._interactiveRuns) {
+      try {
+        await run.kill();
+      } catch (error) {
+        if (env['VERBOSE'] === 'true') {
+          console.warn('Failed to kill interactive run during cleanup:', error);
+        }
+      }
+    }
+    this._interactiveRuns = [];
+
     if (
       process.env['REGENERATE_MODEL_GOLDENS'] === 'true' &&
       this.fakeResponsesPath
@@ -1054,6 +1067,7 @@ export class TestRig {
     const ptyProcess = pty.spawn(executable, commandArgs, ptyOptions);
 
     const run = new InteractiveRun(ptyProcess);
+    this._interactiveRuns.push(run);
     // Wait for the app to be ready
     await run.expectText('  Type your message or @path/to/file', 30000);
     return run;

integration-tests/utf-bom-encoding.test.ts

@@ -4,7 +4,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { describe, it, expect, beforeAll, afterAll } from 'vitest';
+import { describe, it, expect, beforeEach, afterEach } from 'vitest';
 import { writeFileSync, readFileSync } from 'node:fs';
 import { join, resolve } from 'node:path';
 import { TestRig } from './test-helper.js';
@@ -47,28 +47,24 @@ const utf32BE = (s: string) => {
   return Buffer.concat([bom, payload]);
 };
 
-let rig: TestRig;
-let dir: string;
-
 describe('BOM end-to-end integraion', () => {
-  beforeAll(async () => {
+  let rig: TestRig;
+
+  beforeEach(async () => {
     rig = new TestRig();
     await rig.setup('bom-integration', {
       settings: { tools: { core: ['read_file'] } },
     });
-    dir = rig.testDir!;
   });
 
-  afterAll(async () => {
+  afterEach(async () => await rig.cleanup());
-    await rig.cleanup();
-  });
 
   async function runAndAssert(
     filename: string,
     content: Buffer,
     expectedText: string | null,
   ) {
-    writeFileSync(join(dir, filename), content);
+    writeFileSync(join(rig.testDir!, filename), content);
     const prompt = `read the file ${filename} and output its exact contents`;
     const output = await rig.run(prompt);
     await rig.waitForToolCall('read_file');
@@ -128,7 +124,7 @@ describe('BOM end-to-end integraion', () => {
     );
     const imageContent = readFileSync(imagePath);
     const filename = 'gemini-screenshot.png';
-    writeFileSync(join(dir, filename), imageContent);
+    writeFileSync(join(rig.testDir!, filename), imageContent);
     const prompt = `What is shown in the image ${filename}?`;
     const output = await rig.run(prompt);
     await rig.waitForToolCall('read_file');

integration-tests/write_file.test.ts

@@ -4,7 +4,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { describe, it, expect, vi } from 'vitest';
+import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
 import {
   TestRig,
   createToolCallErrorMessage,
@@ -13,8 +13,15 @@ import {
 } from './test-helper.js';
 
 describe('write_file', () => {
+  let rig: TestRig;
+
+  beforeEach(() => {
+    rig = new TestRig();
+  });
+
+  afterEach(async () => await rig.cleanup());
+
   it('should be able to write a file', async () => {
-    const rig = new TestRig();
     await rig.setup('should be able to write a file', {
       settings: { tools: { core: ['write_file', 'read_file'] } },
     });