MediaMetz/gemini-cli
1
0
Fork 0
mirror of https://github.com/google-gemini/gemini-cli.git synced 2026-03-11 14:40:52 -07:00
Code Issues Packages Projects Releases Wiki Activity

fix: use vi.restoreAllMocks and toMatchSnapshot in dialog tests

Browse Source 
- OverageMenuDialog.test.tsx: add afterEach with vi.restoreAllMocks(),
  remove RadioButtonSelect mock, use renderWithProviders + toMatchSnapshot
  for rendering tests and keyboard input for handler tests
- EmptyWalletDialog.test.tsx: same treatment as OverageMenuDialog
- useQuotaAndFallback.test.ts: vi.clearAllMocks() → vi.restoreAllMocks()
  in afterEach to prevent test pollution
This commit is contained in:
Gaurav Ghosh
2026-02-27 07:39:55 -08:00
parent 5cc9644a16
commit e75c7418d4
5 changed files with 196 additions and 178 deletions
Download Patch File Download Diff File Expand all files Collapse all files

133
packages/cli/src/ui/components/EmptyWalletDialog.test.tsx
View File

@@ -4,24 +4,17 @@
* SPDX-License-Identifier: Apache-2.0
*/
import { render } from '../../test-utils/render.js';
import { renderWithProviders } from '../../test-utils/render.js';
import { waitFor } from '../../test-utils/async.js';
import { act } from 'react';
import {
describe,
it,
expect,
vi,
beforeEach,
afterEach,
type Mock,
} from 'vitest';
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { EmptyWalletDialog } from './EmptyWalletDialog.js';
import { RadioButtonSelect } from './shared/RadioButtonSelect.js';
// Mock the child component to make it easier to test the parent
vi.mock('./shared/RadioButtonSelect.js', () => ({
RadioButtonSelect: vi.fn(),
}));
const writeKey = (stdin: { write: (data: string) => void }, key: string) => {
act(() => {
stdin.write(key);
});
};
describe('EmptyWalletDialog', () => {
const mockOnChoice = vi.fn();
@@ -36,8 +29,8 @@ describe('EmptyWalletDialog', () => {
});
describe('rendering', () => {
it('should render with correct menu options when fallback is available', async () => {
const { unmount, waitUntilReady } = render(
it('should match snapshot with fallback available', async () => {
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
<EmptyWalletDialog
failedModel="gemini-2.5-pro"
fallbackModel="gemini-3-flash-preview"
@@ -47,33 +40,12 @@ describe('EmptyWalletDialog', () => {
);
await waitUntilReady();
expect(RadioButtonSelect).toHaveBeenCalledWith(
expect.objectContaining({
items: [
{
label: 'Get AI Credits - Open browser to purchase credits',
value: 'get_credits',
key: 'get_credits',
},
{
label: 'Switch to gemini-3-flash-preview',
value: 'use_fallback',
key: 'use_fallback',
},
{
label: 'Stop - Abort request',
value: 'stop',
key: 'stop',
},
],
}),
undefined,
);
expect(lastFrame()).toMatchSnapshot();
unmount();
});
it('should omit fallback option when fallbackModel is not provided', async () => {
const { unmount, waitUntilReady } = render(
it('should match snapshot without fallback', async () => {
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
<EmptyWalletDialog
failedModel="gemini-2.5-pro"
onChoice={mockOnChoice}
@@ -81,28 +53,12 @@ describe('EmptyWalletDialog', () => {
);
await waitUntilReady();
expect(RadioButtonSelect).toHaveBeenCalledWith(
expect.objectContaining({
items: [
{
label: 'Get AI Credits - Open browser to purchase credits',
value: 'get_credits',
key: 'get_credits',
},
{
label: 'Stop - Abort request',
value: 'stop',
key: 'stop',
},
],
}),
undefined,
);
expect(lastFrame()).toMatchSnapshot();
unmount();
});
it('should display the model name and usage limit message', async () => {
const { lastFrame, unmount, waitUntilReady } = render(
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
<EmptyWalletDialog
failedModel="gemini-2.5-pro"
onChoice={mockOnChoice}
@@ -117,7 +73,7 @@ describe('EmptyWalletDialog', () => {
});
it('should display purchase prompt and credits update notice', async () => {
const { lastFrame, unmount, waitUntilReady } = render(
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
<EmptyWalletDialog
failedModel="gemini-2.5-pro"
onChoice={mockOnChoice}
@@ -134,7 +90,7 @@ describe('EmptyWalletDialog', () => {
});
it('should display reset time when provided', async () => {
const { lastFrame, unmount, waitUntilReady } = render(
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
<EmptyWalletDialog
failedModel="gemini-2.5-pro"
resetTime="3:45 PM"
@@ -150,7 +106,7 @@ describe('EmptyWalletDialog', () => {
});
it('should not display reset time when not provided', async () => {
const { lastFrame, unmount, waitUntilReady } = render(
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
<EmptyWalletDialog
failedModel="gemini-2.5-pro"
onChoice={mockOnChoice}
@@ -164,7 +120,7 @@ describe('EmptyWalletDialog', () => {
});
it('should display slash command hints', async () => {
const { lastFrame, unmount, waitUntilReady } = render(
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
<EmptyWalletDialog
failedModel="gemini-2.5-pro"
onChoice={mockOnChoice}
@@ -182,7 +138,8 @@ describe('EmptyWalletDialog', () => {
describe('onChoice handling', () => {
it('should call onGetCredits and onChoice when get_credits is selected', async () => {
const { unmount, waitUntilReady } = render(
// get_credits is the first item, so just press Enter
const { unmount, stdin, waitUntilReady } = renderWithProviders(
<EmptyWalletDialog
failedModel="gemini-2.5-pro"
onChoice={mockOnChoice}
@@ -191,19 +148,17 @@ describe('EmptyWalletDialog', () => {
);
await waitUntilReady();
const onSelect = (RadioButtonSelect as Mock).mock.calls[0][0].onSelect;
writeKey(stdin, '\r');
await act(async () => {
onSelect('get_credits');
await waitFor(() => {
expect(mockOnGetCredits).toHaveBeenCalled();
expect(mockOnChoice).toHaveBeenCalledWith('get_credits');
});
expect(mockOnGetCredits).toHaveBeenCalled();
expect(mockOnChoice).toHaveBeenCalledWith('get_credits');
unmount();
});
it('should call onChoice without onGetCredits when onGetCredits is not provided', async () => {
const { unmount, waitUntilReady } = render(
const { unmount, stdin, waitUntilReady } = renderWithProviders(
<EmptyWalletDialog
failedModel="gemini-2.5-pro"
onChoice={mockOnChoice}
@@ -211,18 +166,18 @@ describe('EmptyWalletDialog', () => {
);
await waitUntilReady();
const onSelect = (RadioButtonSelect as Mock).mock.calls[0][0].onSelect;
writeKey(stdin, '\r');
await act(async () => {
onSelect('get_credits');
await waitFor(() => {
expect(mockOnChoice).toHaveBeenCalledWith('get_credits');
});
expect(mockOnChoice).toHaveBeenCalledWith('get_credits');
unmount();
});
it('should call onChoice with use_fallback when selected', async () => {
const { unmount, waitUntilReady } = render(
// With fallback: items are [get_credits, use_fallback, stop]
// use_fallback is the second item: Down + Enter
const { unmount, stdin, waitUntilReady } = renderWithProviders(
<EmptyWalletDialog
failedModel="gemini-2.5-pro"
fallbackModel="gemini-3-flash-preview"
@@ -231,17 +186,19 @@ describe('EmptyWalletDialog', () => {
);
await waitUntilReady();
const onSelect = (RadioButtonSelect as Mock).mock.calls[0][0].onSelect;
await act(async () => {
onSelect('use_fallback');
});
writeKey(stdin, '\x1b[B'); // Down arrow
writeKey(stdin, '\r');
expect(mockOnChoice).toHaveBeenCalledWith('use_fallback');
await waitFor(() => {
expect(mockOnChoice).toHaveBeenCalledWith('use_fallback');
});
unmount();
});
it('should call onChoice with stop when selected', async () => {
const { unmount, waitUntilReady } = render(
// Without fallback: items are [get_credits, stop]
// stop is the second item: Down + Enter
const { unmount, stdin, waitUntilReady } = renderWithProviders(
<EmptyWalletDialog
failedModel="gemini-2.5-pro"
onChoice={mockOnChoice}
@@ -249,12 +206,12 @@ describe('EmptyWalletDialog', () => {
);
await waitUntilReady();
const onSelect = (RadioButtonSelect as Mock).mock.calls[0][0].onSelect;
await act(async () => {
onSelect('stop');
});
writeKey(stdin, '\x1b[B'); // Down arrow
writeKey(stdin, '\r');
expect(mockOnChoice).toHaveBeenCalledWith('stop');
await waitFor(() => {
expect(mockOnChoice).toHaveBeenCalledWith('stop');
});
unmount();
});
});

143
packages/cli/src/ui/components/OverageMenuDialog.test.tsx
View File

@@ -4,16 +4,17 @@
* SPDX-License-Identifier: Apache-2.0
*/
import { render } from '../../test-utils/render.js';
import { renderWithProviders } from '../../test-utils/render.js';
import { waitFor } from '../../test-utils/async.js';
import { act } from 'react';
import { describe, it, expect, vi, beforeEach, type Mock } from 'vitest';
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { OverageMenuDialog } from './OverageMenuDialog.js';
import { RadioButtonSelect } from './shared/RadioButtonSelect.js';
// Mock the child component to make it easier to test the parent
vi.mock('./shared/RadioButtonSelect.js', () => ({
RadioButtonSelect: vi.fn(),
}));
const writeKey = (stdin: { write: (data: string) => void }, key: string) => {
act(() => {
stdin.write(key);
});
};
describe('OverageMenuDialog', () => {
const mockOnChoice = vi.fn();
@@ -22,9 +23,13 @@ describe('OverageMenuDialog', () => {
vi.clearAllMocks();
});
afterEach(() => {
vi.restoreAllMocks();
});
describe('rendering', () => {
it('should render with correct menu options when fallback is available', async () => {
const { unmount, waitUntilReady } = render(
it('should match snapshot with fallback available', async () => {
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
<OverageMenuDialog
failedModel="gemini-2.5-pro"
fallbackModel="gemini-3-flash-preview"
@@ -35,38 +40,12 @@ describe('OverageMenuDialog', () => {
);
await waitUntilReady();
expect(RadioButtonSelect).toHaveBeenCalledWith(
expect.objectContaining({
items: [
{
label: 'Use AI Credits - Continue this request (Overage)',
value: 'use_credits',
key: 'use_credits',
},
{
label: 'Manage - View balance and purchase more credits',
value: 'manage',
key: 'manage',
},
{
label: 'Switch to gemini-3-flash-preview',
value: 'use_fallback',
key: 'use_fallback',
},
{
label: 'Stop - Abort request',
value: 'stop',
key: 'stop',
},
],
}),
undefined,
);
expect(lastFrame()).toMatchSnapshot();
unmount();
});
it('should omit fallback option when fallbackModel is not provided', async () => {
const { unmount, waitUntilReady } = render(
it('should match snapshot without fallback', async () => {
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
<OverageMenuDialog
failedModel="gemini-2.5-pro"
creditBalance={500}
@@ -75,33 +54,12 @@ describe('OverageMenuDialog', () => {
);
await waitUntilReady();
expect(RadioButtonSelect).toHaveBeenCalledWith(
expect.objectContaining({
items: [
{
label: 'Use AI Credits - Continue this request (Overage)',
value: 'use_credits',
key: 'use_credits',
},
{
label: 'Manage - View balance and purchase more credits',
value: 'manage',
key: 'manage',
},
{
label: 'Stop - Abort request',
value: 'stop',
key: 'stop',
},
],
}),
undefined,
);
expect(lastFrame()).toMatchSnapshot();
unmount();
});
it('should display the credit balance', async () => {
const { lastFrame, unmount, waitUntilReady } = render(
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
<OverageMenuDialog
failedModel="gemini-2.5-pro"
creditBalance={200}
@@ -117,7 +75,7 @@ describe('OverageMenuDialog', () => {
});
it('should display the model name', async () => {
const { lastFrame, unmount, waitUntilReady } = render(
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
<OverageMenuDialog
failedModel="gemini-2.5-pro"
creditBalance={100}
@@ -133,7 +91,7 @@ describe('OverageMenuDialog', () => {
});
it('should display reset time when provided', async () => {
const { lastFrame, unmount, waitUntilReady } = render(
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
<OverageMenuDialog
failedModel="gemini-2.5-pro"
resetTime="3:45 PM"
@@ -150,7 +108,7 @@ describe('OverageMenuDialog', () => {
});
it('should not display reset time when not provided', async () => {
const { lastFrame, unmount, waitUntilReady } = render(
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
<OverageMenuDialog
failedModel="gemini-2.5-pro"
creditBalance={100}
@@ -165,7 +123,7 @@ describe('OverageMenuDialog', () => {
});
it('should display slash command hints', async () => {
const { lastFrame, unmount, waitUntilReady } = render(
const { lastFrame, unmount, waitUntilReady } = renderWithProviders(
<OverageMenuDialog
failedModel="gemini-2.5-pro"
creditBalance={100}
@@ -184,7 +142,8 @@ describe('OverageMenuDialog', () => {
describe('onChoice handling', () => {
it('should call onChoice with use_credits when selected', async () => {
const { unmount, waitUntilReady } = render(
// use_credits is the first item, so just press Enter
const { unmount, stdin, waitUntilReady } = renderWithProviders(
<OverageMenuDialog
failedModel="gemini-2.5-pro"
creditBalance={100}
@@ -193,17 +152,17 @@ describe('OverageMenuDialog', () => {
);
await waitUntilReady();
const onSelect = (RadioButtonSelect as Mock).mock.calls[0][0].onSelect;
await act(async () => {
onSelect('use_credits');
});
writeKey(stdin, '\r');
expect(mockOnChoice).toHaveBeenCalledWith('use_credits');
await waitFor(() => {
expect(mockOnChoice).toHaveBeenCalledWith('use_credits');
});
unmount();
});
it('should call onChoice with manage when selected', async () => {
const { unmount, waitUntilReady } = render(
// manage is the second item: Down + Enter
const { unmount, stdin, waitUntilReady } = renderWithProviders(
<OverageMenuDialog
failedModel="gemini-2.5-pro"
creditBalance={100}
@@ -212,17 +171,19 @@ describe('OverageMenuDialog', () => {
);
await waitUntilReady();
const onSelect = (RadioButtonSelect as Mock).mock.calls[0][0].onSelect;
await act(async () => {
onSelect('manage');
});
writeKey(stdin, '\x1b[B'); // Down arrow
writeKey(stdin, '\r');
expect(mockOnChoice).toHaveBeenCalledWith('manage');
await waitFor(() => {
expect(mockOnChoice).toHaveBeenCalledWith('manage');
});
unmount();
});
it('should call onChoice with use_fallback when selected', async () => {
const { unmount, waitUntilReady } = render(
// With fallback: items are [use_credits, manage, use_fallback, stop]
// use_fallback is the third item: Down x2 + Enter
const { unmount, stdin, waitUntilReady } = renderWithProviders(
<OverageMenuDialog
failedModel="gemini-2.5-pro"
fallbackModel="gemini-3-flash-preview"
@@ -232,17 +193,20 @@ describe('OverageMenuDialog', () => {
);
await waitUntilReady();
const onSelect = (RadioButtonSelect as Mock).mock.calls[0][0].onSelect;
await act(async () => {
onSelect('use_fallback');
});
writeKey(stdin, '\x1b[B'); // Down arrow
writeKey(stdin, '\x1b[B'); // Down arrow
writeKey(stdin, '\r');
expect(mockOnChoice).toHaveBeenCalledWith('use_fallback');
await waitFor(() => {
expect(mockOnChoice).toHaveBeenCalledWith('use_fallback');
});
unmount();
});
it('should call onChoice with stop when selected', async () => {
const { unmount, waitUntilReady } = render(
// Without fallback: items are [use_credits, manage, stop]
// stop is the third item: Down x2 + Enter
const { unmount, stdin, waitUntilReady } = renderWithProviders(
<OverageMenuDialog
failedModel="gemini-2.5-pro"
creditBalance={100}
@@ -251,12 +215,13 @@ describe('OverageMenuDialog', () => {
);
await waitUntilReady();
const onSelect = (RadioButtonSelect as Mock).mock.calls[0][0].onSelect;
await act(async () => {
onSelect('stop');
});
writeKey(stdin, '\x1b[B'); // Down arrow
writeKey(stdin, '\x1b[B'); // Down arrow
writeKey(stdin, '\r');
expect(mockOnChoice).toHaveBeenCalledWith('stop');
await waitFor(() => {
expect(mockOnChoice).toHaveBeenCalledWith('stop');
});
unmount();
});
});

49
packages/cli/src/ui/components/__snapshots__/EmptyWalletDialog.test.tsx.snap Normal file
View File

@@ -0,0 +1,49 @@
// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
exports[`EmptyWalletDialog > rendering > should match snapshot with fallback available 1`] = `
"╭──────────────────────────────────────────────────────────────────────────────────────────────────╮
│ │
│ Usage limit reached for gemini-2.5-pro. │
│ Access resets at 2:00 PM. │
│ /stats model for usage details │
│ /model to switch models. │
│ /auth to switch to API key. │
│ │
│ To continue using this model now, purchase more AI Credits. │
│ │
│ Newly purchased AI credits may take a few minutes to update. │
│ │
│ How would you like to proceed? │
│ │
│ │
│ ● 1. Get AI Credits - Open browser to purchase credits │
│ 2. Switch to gemini-3-flash-preview │
│ 3. Stop - Abort request │
│ │
│ │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
"
`;
exports[`EmptyWalletDialog > rendering > should match snapshot without fallback 1`] = `
"╭──────────────────────────────────────────────────────────────────────────────────────────────────╮
│ │
│ Usage limit reached for gemini-2.5-pro. │
│ /stats model for usage details │
│ /model to switch models. │
│ /auth to switch to API key. │
│ │
│ To continue using this model now, purchase more AI Credits. │
│ │
│ Newly purchased AI credits may take a few minutes to update. │
│ │
│ How would you like to proceed? │
│ │
│ │
│ ● 1. Get AI Credits - Open browser to purchase credits │
│ 2. Stop - Abort request │
│ │
│ │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
"
`;

47
packages/cli/src/ui/components/__snapshots__/OverageMenuDialog.test.tsx.snap Normal file
View File

@@ -0,0 +1,47 @@
// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
exports[`OverageMenuDialog > rendering > should match snapshot with fallback available 1`] = `
"╭──────────────────────────────────────────────────────────────────────────────────────────────────╮
│ │
│ Usage limit reached for gemini-2.5-pro. │
│ Access resets at 2:00 PM. │
│ /stats model for usage details │
│ /model to switch models. │
│ /auth to switch to API key. │
│ │
│ You have 500 AI Credits available. │
│ │
│ How would you like to proceed? │
│ │
│ │
│ ● 1. Use AI Credits - Continue this request (Overage) │
│ 2. Manage - View balance and purchase more credits │
│ 3. Switch to gemini-3-flash-preview │
│ 4. Stop - Abort request │
│ │
│ │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
"
`;
exports[`OverageMenuDialog > rendering > should match snapshot without fallback 1`] = `
"╭──────────────────────────────────────────────────────────────────────────────────────────────────╮
│ │
│ Usage limit reached for gemini-2.5-pro. │
│ /stats model for usage details │
│ /model to switch models. │
│ /auth to switch to API key. │
│ │
│ You have 500 AI Credits available. │
│ │
│ How would you like to proceed? │
│ │
│ │
│ ● 1. Use AI Credits - Continue this request (Overage) │
│ 2. Manage - View balance and purchase more credits │
│ 3. Stop - Abort request │
│ │
│ │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
"
`;

2
packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts
View File

@@ -103,7 +103,7 @@ describe('useQuotaAndFallback', () => {
});
afterEach(() => {
vi.clearAllMocks();
vi.restoreAllMocks();
});
it('should register a fallback handler on initialization', () => {