Enhance TestRig with process management and timeouts (#15908)

packages/cli/src/utils/cleanup.test.ts

@@ -12,6 +12,8 @@ vi.mock('@google/gemini-cli-core', () => ({
   Storage: vi.fn().mockImplementation(() => ({
     getProjectTempDir: vi.fn().mockReturnValue('/tmp/project'),
   })),
+  shutdownTelemetry: vi.fn(),
+  isTelemetrySdkInitialized: vi.fn().mockReturnValue(false),
 }));
 
 vi.mock('node:fs', () => ({
@@ -20,61 +22,61 @@ vi.mock('node:fs', () => ({
   },
 }));
 
+import {
+  registerCleanup,
+  runExitCleanup,
+  registerSyncCleanup,
+  runSyncCleanup,
+  cleanupCheckpoints,
+  resetCleanupForTesting,
+} from './cleanup.js';
+
 describe('cleanup', () => {
   beforeEach(async () => {
-    vi.resetModules();
     vi.clearAllMocks();
-    // No need to re-assign, we can use the imported functions directly
-    // because we are using vi.resetModules() and re-importing if necessary,
-    // but actually, since we are mocking dependencies, we might not need to re-import cleanup.js
-    // unless it has internal state that needs resetting. It does (cleanupFunctions array).
-    // So we DO need to re-import it to get fresh state.
+    resetCleanupForTesting();
   });
 
   it('should run a registered synchronous function', async () => {
-    const cleanupModule = await import('./cleanup.js');
     const cleanupFn = vi.fn();
-    cleanupModule.registerCleanup(cleanupFn);
+    registerCleanup(cleanupFn);
 
-    await cleanupModule.runExitCleanup();
+    await runExitCleanup();
 
     expect(cleanupFn).toHaveBeenCalledTimes(1);
   });
 
   it('should run a registered asynchronous function', async () => {
-    const cleanupModule = await import('./cleanup.js');
     const cleanupFn = vi.fn().mockResolvedValue(undefined);
-    cleanupModule.registerCleanup(cleanupFn);
+    registerCleanup(cleanupFn);
 
-    await cleanupModule.runExitCleanup();
+    await runExitCleanup();
 
     expect(cleanupFn).toHaveBeenCalledTimes(1);
   });
 
   it('should run multiple registered functions', async () => {
-    const cleanupModule = await import('./cleanup.js');
     const syncFn = vi.fn();
     const asyncFn = vi.fn().mockResolvedValue(undefined);
 
-    cleanupModule.registerCleanup(syncFn);
-    cleanupModule.registerCleanup(asyncFn);
+    registerCleanup(syncFn);
+    registerCleanup(asyncFn);
 
-    await cleanupModule.runExitCleanup();
+    await runExitCleanup();
 
     expect(syncFn).toHaveBeenCalledTimes(1);
     expect(asyncFn).toHaveBeenCalledTimes(1);
   });
 
   it('should continue running cleanup functions even if one throws an error', async () => {
-    const cleanupModule = await import('./cleanup.js');
     const errorFn = vi.fn().mockImplementation(() => {
       throw new Error('test error');
     });
     const successFn = vi.fn();
-    cleanupModule.registerCleanup(errorFn);
-    cleanupModule.registerCleanup(successFn);
+    registerCleanup(errorFn);
+    registerCleanup(successFn);
 
-    await expect(cleanupModule.runExitCleanup()).resolves.not.toThrow();
+    await expect(runExitCleanup()).resolves.not.toThrow();
 
     expect(errorFn).toHaveBeenCalledTimes(1);
     expect(successFn).toHaveBeenCalledTimes(1);
@@ -82,23 +84,21 @@ describe('cleanup', () => {
 
   describe('sync cleanup', () => {
     it('should run registered sync functions', async () => {
-      const cleanupModule = await import('./cleanup.js');
       const syncFn = vi.fn();
-      cleanupModule.registerSyncCleanup(syncFn);
-      cleanupModule.runSyncCleanup();
+      registerSyncCleanup(syncFn);
+      runSyncCleanup();
       expect(syncFn).toHaveBeenCalledTimes(1);
     });
 
     it('should continue running sync cleanup functions even if one throws', async () => {
-      const cleanupModule = await import('./cleanup.js');
       const errorFn = vi.fn().mockImplementation(() => {
         throw new Error('test error');
       });
       const successFn = vi.fn();
-      cleanupModule.registerSyncCleanup(errorFn);
-      cleanupModule.registerSyncCleanup(successFn);
+      registerSyncCleanup(errorFn);
+      registerSyncCleanup(successFn);
 
-      expect(() => cleanupModule.runSyncCleanup()).not.toThrow();
+      expect(() => runSyncCleanup()).not.toThrow();
       expect(errorFn).toHaveBeenCalledTimes(1);
       expect(successFn).toHaveBeenCalledTimes(1);
     });
@@ -106,8 +106,7 @@ describe('cleanup', () => {
 
   describe('cleanupCheckpoints', () => {
     it('should remove checkpoints directory', async () => {
-      const cleanupModule = await import('./cleanup.js');
-      await cleanupModule.cleanupCheckpoints();
+      await cleanupCheckpoints();
       expect(fs.rm).toHaveBeenCalledWith(
         path.join('/tmp/project', 'checkpoints'),
         {
@@ -118,9 +117,8 @@ describe('cleanup', () => {
     });
 
     it('should ignore errors during checkpoint removal', async () => {
-      const cleanupModule = await import('./cleanup.js');
       vi.mocked(fs.rm).mockRejectedValue(new Error('Failed to remove'));
-      await expect(cleanupModule.cleanupCheckpoints()).resolves.not.toThrow();
+      await expect(cleanupCheckpoints()).resolves.not.toThrow();
     });
   });
 });

packages/cli/src/utils/cleanup.ts

@@ -25,6 +25,16 @@ export function registerSyncCleanup(fn: () => void) {
   syncCleanupFunctions.push(fn);
 }
 
+/**
+ * Resets the internal cleanup state for testing purposes.
+ * This allows tests to run in isolation without vi.resetModules().
+ */
+export function resetCleanupForTesting() {
+  cleanupFunctions.length = 0;
+  syncCleanupFunctions.length = 0;
+  configForTelemetry = null;
+}
+
 export function runSyncCleanup() {
   for (const fn of syncCleanupFunctions) {
     try {