feat(autoupdate): Improve update check and refactor for testability (#5389)

1 files changed, 141 insertions, 30 deletions

diff --git a/packages/cli/src/utils/handleAutoUpdate.test.ts b/packages/cli/src/utils/handleAutoUpdate.test.ts
index adaed932..f292d0c2 100644
--- a/packages/cli/src/utils/handleAutoUpdate.test.ts
+++ b/packages/cli/src/utils/handleAutoUpdate.test.ts
@@ -4,22 +4,13 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
-import { ChildProcess, spawn } from 'node:child_process';
-import { handleAutoUpdate } from './handleAutoUpdate.js';
+import { describe, it, expect, vi, beforeEach, afterEach, Mock } from 'vitest';
 import { getInstallationInfo, PackageManager } from './installationInfo.js';
 import { updateEventEmitter } from './updateEventEmitter.js';
 import { UpdateObject } from '../ui/utils/updateCheck.js';
 import { LoadedSettings } from '../config/settings.js';
-
-// Mock dependencies
-vi.mock('node:child_process', async () => {
-  const actual = await vi.importActual('node:child_process');
-  return {
-    ...actual,
-    spawn: vi.fn(),
-  };
-});
+import EventEmitter from 'node:events';
+import { handleAutoUpdate } from './handleAutoUpdate.js';
 
 vi.mock('./installationInfo.js', async () => {
   const actual = await vi.importActual('./installationInfo.js');
@@ -40,21 +31,26 @@ vi.mock('./updateEventEmitter.js', async () => {
   };
 });
 
-const mockSpawn = vi.mocked(spawn);
+interface MockChildProcess extends EventEmitter {
+  stdin: EventEmitter & {
+    write: Mock;
+    end: Mock;
+  };
+  stderr: EventEmitter;
+}
+
 const mockGetInstallationInfo = vi.mocked(getInstallationInfo);
 const mockUpdateEventEmitter = vi.mocked(updateEventEmitter);
 
 describe('handleAutoUpdate', () => {
+  let mockSpawn: Mock;
   let mockUpdateInfo: UpdateObject;
   let mockSettings: LoadedSettings;
-  let mockChildProcess: {
-    stderr: { on: ReturnType<typeof vi.fn> };
-    stdout: { on: ReturnType<typeof vi.fn> };
-    on: ReturnType<typeof vi.fn>;
-    unref: ReturnType<typeof vi.fn>;
-  };
+  let mockChildProcess: MockChildProcess;
 
   beforeEach(() => {
+    mockSpawn = vi.fn();
+    vi.clearAllMocks();
     mockUpdateInfo = {
       update: {
         latest: '2.0.0',
@@ -71,13 +67,17 @@ describe('handleAutoUpdate', () => {
       },
     } as LoadedSettings;
 
-    mockChildProcess = {
-      stdout: { on: vi.fn() },
-      stderr: { on: vi.fn() },
-      on: vi.fn(),
-      unref: vi.fn(),
-    };
-    mockSpawn.mockReturnValue(mockChildProcess as unknown as ChildProcess);
+    mockChildProcess = Object.assign(new EventEmitter(), {
+      stdin: Object.assign(new EventEmitter(), {
+        write: vi.fn(),
+        end: vi.fn(),
+      }),
+      stderr: new EventEmitter(),
+    }) as MockChildProcess;
+
+    mockSpawn.mockReturnValue(
+      mockChildProcess as unknown as ReturnType<typeof mockSpawn>,
+    );
   });
 
   afterEach(() => {
@@ -85,7 +85,7 @@ describe('handleAutoUpdate', () => {
   });
 
   it('should do nothing if update info is null', () => {
-    handleAutoUpdate(null, mockSettings, '/root');
+    handleAutoUpdate(null, mockSettings, '/root', mockSpawn);
     expect(mockGetInstallationInfo).not.toHaveBeenCalled();
     expect(mockUpdateEventEmitter.emit).not.toHaveBeenCalled();
     expect(mockSpawn).not.toHaveBeenCalled();
@@ -100,7 +100,7 @@ describe('handleAutoUpdate', () => {
       packageManager: PackageManager.NPM,
     });
 
-    handleAutoUpdate(mockUpdateInfo, mockSettings, '/root');
+    handleAutoUpdate(mockUpdateInfo, mockSettings, '/root', mockSpawn);
 
     expect(mockUpdateEventEmitter.emit).toHaveBeenCalledTimes(1);
     expect(mockUpdateEventEmitter.emit).toHaveBeenCalledWith(
@@ -120,7 +120,7 @@ describe('handleAutoUpdate', () => {
       packageManager: PackageManager.NPM,
     });
 
-    handleAutoUpdate(mockUpdateInfo, mockSettings, '/root');
+    handleAutoUpdate(mockUpdateInfo, mockSettings, '/root', mockSpawn);
 
     expect(mockUpdateEventEmitter.emit).toHaveBeenCalledTimes(1);
     expect(mockUpdateEventEmitter.emit).toHaveBeenCalledWith(
@@ -140,7 +140,7 @@ describe('handleAutoUpdate', () => {
       packageManager: PackageManager.NPM,
     });
 
-    handleAutoUpdate(mockUpdateInfo, mockSettings, '/root');
+    handleAutoUpdate(mockUpdateInfo, mockSettings, '/root', mockSpawn);
 
     expect(mockUpdateEventEmitter.emit).toHaveBeenCalledTimes(1);
     expect(mockUpdateEventEmitter.emit).toHaveBeenCalledWith(
@@ -150,4 +150,115 @@ describe('handleAutoUpdate', () => {
       },
     );
   });
+
+  it('should attempt to perform an update when conditions are met', async () => {
+    mockGetInstallationInfo.mockReturnValue({
+      updateCommand: 'npm i -g @google/gemini-cli@latest',
+      updateMessage: 'This is an additional message.',
+      isGlobal: false,
+      packageManager: PackageManager.NPM,
+    });
+
+    // Simulate successful execution
+    setTimeout(() => {
+      mockChildProcess.emit('close', 0);
+    }, 0);
+
+    handleAutoUpdate(mockUpdateInfo, mockSettings, '/root', mockSpawn);
+
+    expect(mockSpawn).toHaveBeenCalledOnce();
+  });
+
+  it('should emit "update-failed" when the update process fails', async () => {
+    await new Promise<void>((resolve) => {
+      mockGetInstallationInfo.mockReturnValue({
+        updateCommand: 'npm i -g @google/gemini-cli@latest',
+        updateMessage: 'This is an additional message.',
+        isGlobal: false,
+        packageManager: PackageManager.NPM,
+      });
+
+      // Simulate failed execution
+      setTimeout(() => {
+        mockChildProcess.stderr.emit('data', 'An error occurred');
+        mockChildProcess.emit('close', 1);
+        resolve();
+      }, 0);
+
+      handleAutoUpdate(mockUpdateInfo, mockSettings, '/root', mockSpawn);
+    });
+
+    expect(mockUpdateEventEmitter.emit).toHaveBeenCalledWith('update-failed', {
+      message:
+        'Automatic update failed. Please try updating manually. (command: npm i -g @google/[email protected], stderr: An error occurred)',
+    });
+  });
+
+  it('should emit "update-failed" when the spawn function throws an error', async () => {
+    await new Promise<void>((resolve) => {
+      mockGetInstallationInfo.mockReturnValue({
+        updateCommand: 'npm i -g @google/gemini-cli@latest',
+        updateMessage: 'This is an additional message.',
+        isGlobal: false,
+        packageManager: PackageManager.NPM,
+      });
+
+      // Simulate an error event
+      setTimeout(() => {
+        mockChildProcess.emit('error', new Error('Spawn error'));
+        resolve();
+      }, 0);
+
+      handleAutoUpdate(mockUpdateInfo, mockSettings, '/root', mockSpawn);
+    });
+
+    expect(mockUpdateEventEmitter.emit).toHaveBeenCalledWith('update-failed', {
+      message:
+        'Automatic update failed. Please try updating manually. (error: Spawn error)',
+    });
+  });
+
+  it('should use the "@nightly" tag for nightly updates', async () => {
+    mockUpdateInfo.update.latest = '2.0.0-nightly';
+    mockGetInstallationInfo.mockReturnValue({
+      updateCommand: 'npm i -g @google/gemini-cli@latest',
+      updateMessage: 'This is an additional message.',
+      isGlobal: false,
+      packageManager: PackageManager.NPM,
+    });
+
+    handleAutoUpdate(mockUpdateInfo, mockSettings, '/root', mockSpawn);
+
+    expect(mockSpawn).toHaveBeenCalledWith(
+      'npm i -g @google/gemini-cli@nightly',
+      {
+        shell: true,
+        stdio: 'pipe',
+      },
+    );
+  });
+
+  it('should emit "update-success" when the update process succeeds', async () => {
+    await new Promise<void>((resolve) => {
+      mockGetInstallationInfo.mockReturnValue({
+        updateCommand: 'npm i -g @google/gemini-cli@latest',
+        updateMessage: 'This is an additional message.',
+        isGlobal: false,
+        packageManager: PackageManager.NPM,
+      });
+
+      // Simulate successful execution
+      setTimeout(() => {
+        mockChildProcess.emit('close', 0);
+        resolve();
+      }, 0);
+
+      handleAutoUpdate(mockUpdateInfo, mockSettings, '/root', mockSpawn);
+    });
+
+    expect(mockUpdateEventEmitter.emit).toHaveBeenCalledWith('update-success', {
+      message:
+        'Update successful! The new version will be used on your next run.',
+    });
+  });
 });