From 678e947f21252f5e948ff9e67d6c19ebed7f5c94 Mon Sep 17 00:00:00 2001
From: Arya Gummadi <aryagummadi@google.com>
Date: Wed, 17 Sep 2025 10:12:33 -0700
Subject: [PATCH] fix(tests): improve test reliability and performance (#8395)

---
 packages/core/src/utils/fileUtils.test.ts     | 37 ++++++++++++-------
 .../core/src/utils/memoryDiscovery.test.ts    | 20 +++++++---
 2 files changed, 37 insertions(+), 20 deletions(-)

diff --git a/packages/core/src/utils/fileUtils.test.ts b/packages/core/src/utils/fileUtils.test.ts
index dd1ad6e62c..2fd4174ab0 100644
--- a/packages/core/src/utils/fileUtils.test.ts
+++ b/packages/core/src/utils/fileUtils.test.ts
@@ -15,6 +15,7 @@ import {
 } from 'vitest';
 
 import * as actualNodeFs from 'node:fs'; // For setup/teardown
+import fs from 'node:fs';
 import fsPromises from 'node:fs/promises';
 import path from 'node:path';
 import os from 'node:os';
@@ -954,22 +955,30 @@ describe('fileUtils', () => {
     });
 
     it('should return an error if the file size exceeds 20MB', async () => {
-      // Create a file just over 20MB
-      const twentyOneMB = 21 * 1024 * 1024;
-      const buffer = Buffer.alloc(twentyOneMB, 0x61); // Fill with 'a'
-      actualNodeFs.writeFileSync(testTextFilePath, buffer);
+      // Create a small test file
+      actualNodeFs.writeFileSync(testTextFilePath, 'test content');
 
-      const result = await processSingleFileContent(
-        testTextFilePath,
-        tempRootDir,
-        new StandardFileSystemService(),
-      );
+      // Spy on fs.promises.stat to return a large file size
+      const statSpy = vi.spyOn(fs.promises, 'stat').mockResolvedValueOnce({
+        size: 21 * 1024 * 1024,
+        isDirectory: () => false,
+      } as fs.Stats);
 
-      expect(result.error).toContain('File size exceeds the 20MB limit');
-      expect(result.returnDisplay).toContain(
-        'File size exceeds the 20MB limit',
-      );
-      expect(result.llmContent).toContain('File size exceeds the 20MB limit');
+      try {
+        const result = await processSingleFileContent(
+          testTextFilePath,
+          tempRootDir,
+          new StandardFileSystemService(),
+        );
+
+        expect(result.error).toContain('File size exceeds the 20MB limit');
+        expect(result.returnDisplay).toContain(
+          'File size exceeds the 20MB limit',
+        );
+        expect(result.llmContent).toContain('File size exceeds the 20MB limit');
+      } finally {
+        statSpy.mockRestore();
+      }
     });
   });
 });
diff --git a/packages/core/src/utils/memoryDiscovery.test.ts b/packages/core/src/utils/memoryDiscovery.test.ts
index 67b6d4936e..9cf471b1be 100644
--- a/packages/core/src/utils/memoryDiscovery.test.ts
+++ b/packages/core/src/utils/memoryDiscovery.test.ts
@@ -63,7 +63,13 @@ describe('loadServerHierarchicalMemory', () => {
     // Some tests set this to a different value.
     setGeminiMdFilename(DEFAULT_CONTEXT_FILENAME);
     // Clean up the temporary directory to prevent resource leaks.
-    await fsPromises.rm(testRootDir, { recursive: true, force: true });
+    // Use maxRetries option for robust cleanup without race conditions
+    await fsPromises.rm(testRootDir, {
+      recursive: true,
+      force: true,
+      maxRetries: 3,
+      retryDelay: 10,
+    });
   });
 
   describe('when untrusted', () => {
@@ -354,9 +360,11 @@ describe('loadServerHierarchicalMemory', () => {
       .spyOn(console, 'debug')
       .mockImplementation(() => {});
 
-    for (let i = 0; i < 100; i++) {
-      await createEmptyDir(path.join(cwd, `deep_dir_${i}`));
-    }
+    // Create directories in parallel for better performance
+    const dirPromises = Array.from({ length: 2 }, (_, i) =>
+      createEmptyDir(path.join(cwd, `deep_dir_${i}`)),
+    );
+    await Promise.all(dirPromises);
 
     // Pass the custom limit directly to the function
     await loadServerHierarchicalMemory(
@@ -371,12 +379,12 @@ describe('loadServerHierarchicalMemory', () => {
         respectGitIgnore: true,
         respectGeminiIgnore: true,
       },
-      50, // maxDirs
+      1, // maxDirs
     );
 
     expect(consoleDebugSpy).toHaveBeenCalledWith(
       expect.stringContaining('[DEBUG] [BfsFileSearch]'),
-      expect.stringContaining('Scanning [50/50]:'),
+      expect.stringContaining('Scanning [1/1]:'),
     );
 
     vi.mocked(console.debug).mockRestore();