gemini-cli/packages/cli/src/ui/commands/statsCommand.test.ts

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

import { vi, describe, it, expect, beforeEach } from 'vitest';
import { statsCommand } from './statsCommand.js';
import { type CommandContext } from './types.js';
import { createMockCommandContext } from '../../test-utils/mockCommandContext.js';
import { MessageType } from '../types.js';
import { formatDuration } from '../utils/formatters.js';
import type { Config } from '@google/gemini-cli-core';

vi.mock('@google/gemini-cli-core', async (importOriginal) => {
  const actual =
    await importOriginal<typeof import('@google/gemini-cli-core')>();
  return {
    ...actual,
    UserAccountManager: vi.fn().mockImplementation(() => ({
      getCachedGoogleAccount: vi.fn().mockReturnValue('mock@example.com'),
    })),
    getG1CreditBalance: vi.fn().mockReturnValue(undefined),
  };
});

describe('statsCommand', () => {
  let mockContext: CommandContext;
  const startTime = new Date('2025-07-14T10:00:00.000Z');
  const endTime = new Date('2025-07-14T10:00:30.000Z');

  beforeEach(() => {
    vi.useFakeTimers();
    vi.setSystemTime(endTime);

    // 1. Create the mock context with all default values
    mockContext = createMockCommandContext();

    // 2. Directly set the property on the created mock context
    mockContext.session.stats.sessionStartTime = startTime;
  });

  it('should display general session stats when run with no subcommand', async () => {
    if (!statsCommand.action) throw new Error('Command has no action');

    mockContext.services.config = {
      refreshUserQuota: vi.fn(),
      refreshAvailableCredits: vi.fn(),
      getUserTierName: vi.fn(),
      getUserPaidTier: vi.fn(),
      getModel: vi.fn(),
    } as unknown as Config;

    await statsCommand.action(mockContext, '');

    const expectedDuration = formatDuration(
      endTime.getTime() - startTime.getTime(),
    );
    expect(mockContext.ui.addItem).toHaveBeenCalledWith({
      type: MessageType.STATS,
      duration: expectedDuration,
      selectedAuthType: '',
      tier: undefined,
      userEmail: 'mock@example.com',
      currentModel: undefined,
      creditBalance: undefined,
    });
  });

  it('should fetch and display quota if config is available', async () => {
    if (!statsCommand.action) throw new Error('Command has no action');

    const mockQuota = { buckets: [] };
    const mockRefreshUserQuota = vi.fn().mockResolvedValue(mockQuota);
    const mockGetUserTierName = vi.fn().mockReturnValue('Basic');
    const mockGetModel = vi.fn().mockReturnValue('gemini-pro');
    const mockGetQuotaRemaining = vi.fn().mockReturnValue(85);
    const mockGetQuotaLimit = vi.fn().mockReturnValue(100);
    const mockGetQuotaResetTime = vi
      .fn()
      .mockReturnValue('2025-01-01T12:00:00Z');

    mockContext.services.config = {
      refreshUserQuota: mockRefreshUserQuota,
      getUserTierName: mockGetUserTierName,
      getModel: mockGetModel,
      getQuotaRemaining: mockGetQuotaRemaining,
      getQuotaLimit: mockGetQuotaLimit,
      getQuotaResetTime: mockGetQuotaResetTime,
      getUserPaidTier: vi.fn(),
      refreshAvailableCredits: vi.fn(),
    } as unknown as Config;

    await statsCommand.action(mockContext, '');

    expect(mockRefreshUserQuota).toHaveBeenCalled();
    expect(mockContext.ui.addItem).toHaveBeenCalledWith(
      expect.objectContaining({
        quotas: mockQuota,
        tier: 'Basic',
        currentModel: 'gemini-pro',
        pooledRemaining: 85,
        pooledLimit: 100,
        pooledResetTime: '2025-01-01T12:00:00Z',
      }),
    );
  });

  it('should display model stats when using the "model" subcommand', () => {
    const modelSubCommand = statsCommand.subCommands?.find(
      (sc) => sc.name === 'model',
    );
    if (!modelSubCommand?.action) throw new Error('Subcommand has no action');

    // eslint-disable-next-line @typescript-eslint/no-floating-promises
    modelSubCommand.action(mockContext, '');

    expect(mockContext.ui.addItem).toHaveBeenCalledWith({
      type: MessageType.MODEL_STATS,
      selectedAuthType: '',
      tier: undefined,
      userEmail: 'mock@example.com',
      currentModel: undefined,
      pooledRemaining: undefined,
      pooledLimit: undefined,
    });
  });

  it('should display tool stats when using the "tools" subcommand', () => {
    const toolsSubCommand = statsCommand.subCommands?.find(
      (sc) => sc.name === 'tools',
    );
    if (!toolsSubCommand?.action) throw new Error('Subcommand has no action');

    // eslint-disable-next-line @typescript-eslint/no-floating-promises
    toolsSubCommand.action(mockContext, '');

    expect(mockContext.ui.addItem).toHaveBeenCalledWith({
      type: MessageType.TOOL_STATS,
    });
  });
});