Improve code coverage for cli package (#13724)

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

@@ -5,61 +5,122 @@
  */
 
 import { vi, describe, it, expect, beforeEach } from 'vitest';
-import type { registerCleanup, runExitCleanup } from './cleanup.js';
+import { promises as fs } from 'node:fs';
+import * as path from 'node:path';
+
+vi.mock('@google/gemini-cli-core', () => ({
+  Storage: vi.fn().mockImplementation(() => ({
+    getProjectTempDir: vi.fn().mockReturnValue('/tmp/project'),
+  })),
+}));
+
+vi.mock('node:fs', () => ({
+  promises: {
+    rm: vi.fn(),
+  },
+}));
 
 describe('cleanup', () => {
-  let register: typeof registerCleanup;
-  let runExit: typeof runExitCleanup;
-
   beforeEach(async () => {
     vi.resetModules();
-    const cleanupModule = await import('./cleanup.js');
-    register = cleanupModule.registerCleanup;
-    runExit = cleanupModule.runExitCleanup;
-  }, 30000);
+    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.
+  });
 
   it('should run a registered synchronous function', async () => {
+    const cleanupModule = await import('./cleanup.js');
     const cleanupFn = vi.fn();
-    register(cleanupFn);
+    cleanupModule.registerCleanup(cleanupFn);
 
-    await runExit();
+    await cleanupModule.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);
-    register(cleanupFn);
+    cleanupModule.registerCleanup(cleanupFn);
 
-    await runExit();
+    await cleanupModule.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);
 
-    register(syncFn);
-    register(asyncFn);
+    cleanupModule.registerCleanup(syncFn);
+    cleanupModule.registerCleanup(asyncFn);
 
-    await runExit();
+    await cleanupModule.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();
-    register(errorFn);
-    register(successFn);
+    cleanupModule.registerCleanup(errorFn);
+    cleanupModule.registerCleanup(successFn);
 
-    await expect(runExit()).resolves.not.toThrow();
+    await expect(cleanupModule.runExitCleanup()).resolves.not.toThrow();
 
     expect(errorFn).toHaveBeenCalledTimes(1);
     expect(successFn).toHaveBeenCalledTimes(1);
   });
+
+  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();
+      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);
+
+      expect(() => cleanupModule.runSyncCleanup()).not.toThrow();
+      expect(errorFn).toHaveBeenCalledTimes(1);
+      expect(successFn).toHaveBeenCalledTimes(1);
+    });
+  });
+
+  describe('cleanupCheckpoints', () => {
+    it('should remove checkpoints directory', async () => {
+      const cleanupModule = await import('./cleanup.js');
+      await cleanupModule.cleanupCheckpoints();
+      expect(fs.rm).toHaveBeenCalledWith(
+        path.join('/tmp/project', 'checkpoints'),
+        {
+          recursive: true,
+          force: true,
+        },
+      );
+    });
+
+    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();
+    });
+  });
 });