From e75c7418d409f842fc0884df4e3b3ae2b2535564 Mon Sep 17 00:00:00 2001 From: Gaurav Ghosh Date: Fri, 27 Feb 2026 07:39:55 -0800 Subject: [PATCH] fix: use vi.restoreAllMocks and toMatchSnapshot in dialog tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 --- .../ui/components/EmptyWalletDialog.test.tsx | 133 ++++++---------- .../ui/components/OverageMenuDialog.test.tsx | 143 +++++++----------- .../EmptyWalletDialog.test.tsx.snap | 49 ++++++ .../OverageMenuDialog.test.tsx.snap | 47 ++++++ .../src/ui/hooks/useQuotaAndFallback.test.ts | 2 +- 5 files changed, 196 insertions(+), 178 deletions(-) create mode 100644 packages/cli/src/ui/components/__snapshots__/EmptyWalletDialog.test.tsx.snap create mode 100644 packages/cli/src/ui/components/__snapshots__/OverageMenuDialog.test.tsx.snap diff --git a/packages/cli/src/ui/components/EmptyWalletDialog.test.tsx b/packages/cli/src/ui/components/EmptyWalletDialog.test.tsx index 8dfd3ec2f3..6f8f063c43 100644 --- a/packages/cli/src/ui/components/EmptyWalletDialog.test.tsx +++ b/packages/cli/src/ui/components/EmptyWalletDialog.test.tsx @@ -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( { ); 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( { ); 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( { }); it('should display purchase prompt and credits update notice', async () => { - const { lastFrame, unmount, waitUntilReady } = render( + const { lastFrame, unmount, waitUntilReady } = renderWithProviders( { }); it('should display reset time when provided', async () => { - const { lastFrame, unmount, waitUntilReady } = render( + const { lastFrame, unmount, waitUntilReady } = renderWithProviders( { }); it('should not display reset time when not provided', async () => { - const { lastFrame, unmount, waitUntilReady } = render( + const { lastFrame, unmount, waitUntilReady } = renderWithProviders( { }); it('should display slash command hints', async () => { - const { lastFrame, unmount, waitUntilReady } = render( + const { lastFrame, unmount, waitUntilReady } = renderWithProviders( { 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( { ); 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( { ); 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( { ); 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( { ); 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(); }); }); diff --git a/packages/cli/src/ui/components/OverageMenuDialog.test.tsx b/packages/cli/src/ui/components/OverageMenuDialog.test.tsx index beae41da26..8fb2f7469b 100644 --- a/packages/cli/src/ui/components/OverageMenuDialog.test.tsx +++ b/packages/cli/src/ui/components/OverageMenuDialog.test.tsx @@ -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( { ); 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( { ); 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( { }); it('should display the model name', async () => { - const { lastFrame, unmount, waitUntilReady } = render( + const { lastFrame, unmount, waitUntilReady } = renderWithProviders( { }); it('should display reset time when provided', async () => { - const { lastFrame, unmount, waitUntilReady } = render( + const { lastFrame, unmount, waitUntilReady } = renderWithProviders( { }); it('should not display reset time when not provided', async () => { - const { lastFrame, unmount, waitUntilReady } = render( + const { lastFrame, unmount, waitUntilReady } = renderWithProviders( { }); it('should display slash command hints', async () => { - const { lastFrame, unmount, waitUntilReady } = render( + const { lastFrame, unmount, waitUntilReady } = renderWithProviders( { 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( { ); 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( { ); 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( { ); 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( { ); 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(); }); }); diff --git a/packages/cli/src/ui/components/__snapshots__/EmptyWalletDialog.test.tsx.snap b/packages/cli/src/ui/components/__snapshots__/EmptyWalletDialog.test.tsx.snap new file mode 100644 index 0000000000..c8ebd612af --- /dev/null +++ b/packages/cli/src/ui/components/__snapshots__/EmptyWalletDialog.test.tsx.snap @@ -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 │ +│ │ +│ │ +╰──────────────────────────────────────────────────────────────────────────────────────────────────╯ +" +`; diff --git a/packages/cli/src/ui/components/__snapshots__/OverageMenuDialog.test.tsx.snap b/packages/cli/src/ui/components/__snapshots__/OverageMenuDialog.test.tsx.snap new file mode 100644 index 0000000000..978c5cd17c --- /dev/null +++ b/packages/cli/src/ui/components/__snapshots__/OverageMenuDialog.test.tsx.snap @@ -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 │ +│ │ +│ │ +╰──────────────────────────────────────────────────────────────────────────────────────────────────╯ +" +`; diff --git a/packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts b/packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts index 52caa38f9d..b71233df68 100644 --- a/packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts +++ b/packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts @@ -103,7 +103,7 @@ describe('useQuotaAndFallback', () => { }); afterEach(() => { - vi.clearAllMocks(); + vi.restoreAllMocks(); }); it('should register a fallback handler on initialization', () => {