packages/cli/src/config/extensions/update.test.ts

/**
 * @license
 * Copyright 2025 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */

import { describe, it, expect, vi, beforeEach } from 'vitest';
import {
  updateExtension,
  updateAllUpdatableExtensions,
  checkForAllExtensionUpdates,
} from './update.js';
import {
  ExtensionUpdateState,
  type ExtensionUpdateStatus,
} from '../../ui/state/extensions.js';
import { ExtensionStorage } from './storage.js';
import { type ExtensionManager, copyExtension } from '../extension-manager.js';
import { checkForExtensionUpdate } from './github.js';
import { loadInstallMetadata } from '../extension.js';
import * as fs from 'node:fs';
import {
  type GeminiCLIExtension,
  type ExtensionInstallMetadata,
  IntegrityDataStatus,
} from '@google/gemini-cli-core';

vi.mock('./storage.js', () => ({
  ExtensionStorage: {
    createTmpDir: vi.fn(),
  },
}));

vi.mock('../extension-manager.js', () => ({
  copyExtension: vi.fn(),
  // We don't need to mock the class implementation if we pass a mock instance
}));

vi.mock('./github.js', () => ({
  checkForExtensionUpdate: vi.fn(),
}));

vi.mock('../extension.js', () => ({
  loadInstallMetadata: vi.fn(),
}));

vi.mock('node:fs', async (importOriginal) => {
  const actual = await importOriginal<typeof import('node:fs')>();
  return {
    ...actual,
    promises: {
      ...actual.promises,
      rm: vi.fn(),
    },
  };
});

describe('Extension Update Logic', () => {
  let mockExtensionManager: ExtensionManager;
  let mockDispatch: ReturnType<typeof vi.fn>;
  const mockExtension: GeminiCLIExtension = {
    name: 'test-extension',
    version: '1.0.0',
    path: '/path/to/extension',
  } as GeminiCLIExtension;

  beforeEach(() => {
    vi.clearAllMocks();
    mockExtensionManager = {
      loadExtensionConfig: vi.fn().mockResolvedValue({
        name: 'test-extension',
        version: '1.0.0',
      }),
      installOrUpdateExtension: vi.fn().mockResolvedValue({
        ...mockExtension,
        version: '1.1.0',
      }),
      verifyExtensionIntegrity: vi
        .fn()
        .mockResolvedValue(IntegrityDataStatus.VERIFIED),
      storeExtensionIntegrity: vi.fn().mockResolvedValue(undefined),
    } as unknown as ExtensionManager;
    mockDispatch = vi.fn();

    // Default mock behaviors
    vi.mocked(ExtensionStorage.createTmpDir).mockResolvedValue('/tmp/mock-dir');
    vi.mocked(loadInstallMetadata).mockReturnValue({
      source: 'https://example.com/repo.git',
      type: 'git',
    });
  });

  describe('updateExtension', () => {
    it('should return undefined if state is already UPDATING', async () => {
      const result = await updateExtension(
        mockExtension,
        mockExtensionManager,
        ExtensionUpdateState.UPDATING,
        mockDispatch,
      );
      expect(result).toBeUndefined();
      expect(mockDispatch).not.toHaveBeenCalled();
    });

    it('should throw error and set state to ERROR if install metadata type is unknown', async () => {
      vi.mocked(loadInstallMetadata).mockReturnValue({
        type: undefined,
      } as unknown as ExtensionInstallMetadata);

      await expect(
        updateExtension(
          mockExtension,
          mockExtensionManager,
          ExtensionUpdateState.UPDATE_AVAILABLE,
          mockDispatch,
        ),
      ).rejects.toThrow('type is unknown');

      expect(mockDispatch).toHaveBeenCalledWith({
        type: 'SET_STATE',
        payload: {
          name: mockExtension.name,
          state: ExtensionUpdateState.UPDATING,
        },
      });
      expect(mockDispatch).toHaveBeenCalledWith({
        type: 'SET_STATE',
        payload: {
          name: mockExtension.name,
          state: ExtensionUpdateState.ERROR,
        },
      });
    });

    it('should throw error and set state to UP_TO_DATE if extension is linked', async () => {
      vi.mocked(loadInstallMetadata).mockReturnValue({
        type: 'link',
        source: '',
      });

      await expect(
        updateExtension(
          mockExtension,
          mockExtensionManager,
          ExtensionUpdateState.UPDATE_AVAILABLE,
          mockDispatch,
        ),
      ).rejects.toThrow('Extension is linked');

      expect(mockDispatch).toHaveBeenCalledWith({
        type: 'SET_STATE',
        payload: {
          name: mockExtension.name,
          state: ExtensionUpdateState.UP_TO_DATE,
        },
      });
    });

    it('should successfully update extension and set state to UPDATED_NEEDS_RESTART by default', async () => {
      vi.mocked(mockExtensionManager.loadExtensionConfig).mockReturnValue(
        Promise.resolve({
          name: 'test-extension',
          version: '1.0.0',
        }),
      );
      vi.mocked(
        mockExtensionManager.installOrUpdateExtension,
      ).mockResolvedValue({
        ...mockExtension,
        version: '1.1.0',
      });

      const result = await updateExtension(
        mockExtension,
        mockExtensionManager,
        ExtensionUpdateState.UPDATE_AVAILABLE,
        mockDispatch,
      );

      expect(mockExtensionManager.installOrUpdateExtension).toHaveBeenCalled();
      expect(mockDispatch).toHaveBeenCalledWith({
        type: 'SET_STATE',
        payload: {
          name: mockExtension.name,
          state: ExtensionUpdateState.UPDATED_NEEDS_RESTART,
        },
      });
      expect(result).toEqual({
        name: 'test-extension',
        originalVersion: '1.0.0',
        updatedVersion: '1.1.0',
      });
      expect(fs.promises.rm).toHaveBeenCalledWith('/tmp/mock-dir', {
        recursive: true,
        force: true,
      });
    });

    it('should migrate source if migratedTo is set and an update is available', async () => {
      vi.mocked(mockExtensionManager.loadExtensionConfig).mockReturnValue(
        Promise.resolve({
          name: 'test-extension',
          version: '1.0.0',
        }),
      );
      vi.mocked(
        mockExtensionManager.installOrUpdateExtension,
      ).mockResolvedValue({
        ...mockExtension,
        version: '1.1.0',
      });
      vi.mocked(checkForExtensionUpdate).mockResolvedValue(
        ExtensionUpdateState.UPDATE_AVAILABLE,
      );

      const extensionWithMigratedTo = {
        ...mockExtension,
        migratedTo: 'https://new-source.com/repo.git',
      };

      await updateExtension(
        extensionWithMigratedTo,
        mockExtensionManager,
        ExtensionUpdateState.UPDATE_AVAILABLE,
        mockDispatch,
      );

      expect(checkForExtensionUpdate).toHaveBeenCalledWith(
        expect.objectContaining({
          installMetadata: expect.objectContaining({
            source: 'https://new-source.com/repo.git',
          }),
        }),
        mockExtensionManager,
      );

      expect(
        mockExtensionManager.installOrUpdateExtension,
      ).toHaveBeenCalledWith(
        expect.objectContaining({
          source: 'https://new-source.com/repo.git',
        }),
        expect.anything(),
      );
    });

    it('should set state to UPDATED if enableExtensionReloading is true', async () => {
      vi.mocked(mockExtensionManager.loadExtensionConfig).mockReturnValue(
        Promise.resolve({
          name: 'test-extension',
          version: '1.0.0',
        }),
      );
      vi.mocked(
        mockExtensionManager.installOrUpdateExtension,
      ).mockResolvedValue({
        ...mockExtension,
        version: '1.1.0',
      });

      await updateExtension(
        mockExtension,
        mockExtensionManager,
        ExtensionUpdateState.UPDATE_AVAILABLE,
        mockDispatch,
        true, // enableExtensionReloading
      );

      expect(mockDispatch).toHaveBeenCalledWith({
        type: 'SET_STATE',
        payload: {
          name: mockExtension.name,
          state: ExtensionUpdateState.UPDATED,
        },
      });
    });

    it('should rollback and set state to ERROR if installation fails', async () => {
      vi.mocked(mockExtensionManager.loadExtensionConfig).mockReturnValue(
        Promise.resolve({
          name: 'test-extension',
          version: '1.0.0',
        }),
      );
      vi.mocked(
        mockExtensionManager.installOrUpdateExtension,
      ).mockRejectedValue(new Error('Install failed'));

      await expect(
        updateExtension(
          mockExtension,
          mockExtensionManager,
          ExtensionUpdateState.UPDATE_AVAILABLE,
          mockDispatch,
        ),
      ).rejects.toThrow('Updated extension not found after installation');

      expect(copyExtension).toHaveBeenCalledWith(
        '/tmp/mock-dir',
        mockExtension.path,
      );
      expect(mockDispatch).toHaveBeenCalledWith({
        type: 'SET_STATE',
        payload: {
          name: mockExtension.name,
          state: ExtensionUpdateState.ERROR,
        },
      });
      expect(fs.promises.rm).toHaveBeenCalled();
    });

    describe('Integrity Verification', () => {
      it('should fail update with security alert if integrity is invalid', async () => {
        vi.mocked(
          mockExtensionManager.verifyExtensionIntegrity,
        ).mockResolvedValue(IntegrityDataStatus.INVALID);

        await expect(
          updateExtension(
            mockExtension,
            mockExtensionManager,
            ExtensionUpdateState.UPDATE_AVAILABLE,
            mockDispatch,
          ),
        ).rejects.toThrow(
          'Extension test-extension cannot be updated. Extension integrity cannot be verified.',
        );

        expect(mockDispatch).toHaveBeenCalledWith({
          type: 'SET_STATE',
          payload: {
            name: mockExtension.name,
            state: ExtensionUpdateState.ERROR,
          },
        });
      });

      it('should establish trust on first update if integrity data is missing', async () => {
        vi.mocked(
          mockExtensionManager.verifyExtensionIntegrity,
        ).mockResolvedValue(IntegrityDataStatus.MISSING);

        await updateExtension(
          mockExtension,
          mockExtensionManager,
          ExtensionUpdateState.UPDATE_AVAILABLE,
          mockDispatch,
        );

        // Verify updateExtension delegates to installOrUpdateExtension,
        // which is responsible for establishing trust internally.
        expect(
          mockExtensionManager.installOrUpdateExtension,
        ).toHaveBeenCalled();

        expect(mockDispatch).toHaveBeenCalledWith({
          type: 'SET_STATE',
          payload: {
            name: mockExtension.name,
            state: ExtensionUpdateState.UPDATED_NEEDS_RESTART,
          },
        });
      });

      it('should throw if integrity manager throws', async () => {
        vi.mocked(
          mockExtensionManager.verifyExtensionIntegrity,
        ).mockRejectedValue(new Error('Verification failed'));

        await expect(
          updateExtension(
            mockExtension,
            mockExtensionManager,
            ExtensionUpdateState.UPDATE_AVAILABLE,
            mockDispatch,
          ),
        ).rejects.toThrow(
          'Extension test-extension cannot be updated. Verification failed',
        );
      });
    });
  });

  describe('updateAllUpdatableExtensions', () => {
    it('should update all extensions with UPDATE_AVAILABLE status', async () => {
      const extensions: GeminiCLIExtension[] = [
        { ...mockExtension, name: 'ext1' },
        { ...mockExtension, name: 'ext2' },
        { ...mockExtension, name: 'ext3' },
      ];
      const extensionsState = new Map([
        ['ext1', { status: ExtensionUpdateState.UPDATE_AVAILABLE }],
        ['ext2', { status: ExtensionUpdateState.UP_TO_DATE }],
        ['ext3', { status: ExtensionUpdateState.UPDATE_AVAILABLE }],
      ]);

      vi.mocked(mockExtensionManager.loadExtensionConfig).mockReturnValue(
        Promise.resolve({
          name: 'ext',
          version: '1.0.0',
        }),
      );
      vi.mocked(
        mockExtensionManager.installOrUpdateExtension,
      ).mockResolvedValue({ ...mockExtension, version: '1.1.0' });

      const results = await updateAllUpdatableExtensions(
        extensions,
        extensionsState as Map<string, ExtensionUpdateStatus>,
        mockExtensionManager,
        mockDispatch,
      );

      expect(results).toHaveLength(2);
      expect(results.map((r) => r.name)).toEqual(['ext1', 'ext3']);
      expect(
        mockExtensionManager.installOrUpdateExtension,
      ).toHaveBeenCalledTimes(2);
    });
  });

  describe('checkForAllExtensionUpdates', () => {
    it('should dispatch BATCH_CHECK_START and BATCH_CHECK_END', async () => {
      await checkForAllExtensionUpdates([], mockExtensionManager, mockDispatch);

      expect(mockDispatch).toHaveBeenCalledWith({ type: 'BATCH_CHECK_START' });
      expect(mockDispatch).toHaveBeenCalledWith({ type: 'BATCH_CHECK_END' });
    });

    it('should set state to NOT_UPDATABLE if no install metadata', async () => {
      const extensions: GeminiCLIExtension[] = [
        { ...mockExtension, installMetadata: undefined },
      ];

      await checkForAllExtensionUpdates(
        extensions,
        mockExtensionManager,
        mockDispatch,
      );

      expect(mockDispatch).toHaveBeenCalledWith({
        type: 'SET_STATE',
        payload: {
          name: mockExtension.name,
          state: ExtensionUpdateState.NOT_UPDATABLE,
        },
      });
    });

    it('should check for updates and update state', async () => {
      const extensions: GeminiCLIExtension[] = [
        { ...mockExtension, installMetadata: { type: 'git', source: '...' } },
      ];
      vi.mocked(checkForExtensionUpdate).mockResolvedValue(
        ExtensionUpdateState.UPDATE_AVAILABLE,
      );

      await checkForAllExtensionUpdates(
        extensions,
        mockExtensionManager,
        mockDispatch,
      );

      expect(mockDispatch).toHaveBeenCalledWith({
        type: 'SET_STATE',
        payload: {
          name: mockExtension.name,
          state: ExtensionUpdateState.CHECKING_FOR_UPDATES,
        },
      });
      expect(mockDispatch).toHaveBeenCalledWith({
        type: 'SET_STATE',
        payload: {
          name: mockExtension.name,
          state: ExtensionUpdateState.UPDATE_AVAILABLE,
        },
      });
    });
  });
});