diff --git a/integration-tests/replace.test.ts b/integration-tests/replace.test.ts
index a04025e7..e7ed23a2 100644
--- a/integration-tests/replace.test.ts
+++ b/integration-tests/replace.test.ts
@@ -4,7 +4,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { describe, it, expect } from 'vitest';
+import { describe, it, expect, vi } from 'vitest';
 import { TestRig, printDebugInfo, validateModelOutput } from './test-helper.js';
 
 describe('replace', () => {
diff --git a/integration-tests/write_file.test.ts b/integration-tests/write_file.test.ts
index 49cdda38..7dd6445d 100644
--- a/integration-tests/write_file.test.ts
+++ b/integration-tests/write_file.test.ts
@@ -4,7 +4,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { describe, it, expect } from 'vitest';
+import { describe, it, expect, vi } from 'vitest';
 import {
   TestRig,
   createToolCallErrorMessage,