feat: add Pro Quota Dialog (#7094)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Jacob Richman <jacob314@gmail.com>
2026-04-29 14:34:55 -07:00 · 2025-08-29 04:19:43 +05:30
parent a0fbe000ee
commit a63e67823d
3 changed files with 218 additions and 3 deletions
@@ -0,0 +1,89 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { render } from 'ink-testing-library';
+import { describe, it, expect, vi, beforeEach, type Mock } from 'vitest';
+import { ProQuotaDialog } from './ProQuotaDialog.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(),
+}));
+
+describe('ProQuotaDialog', () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  it('should render with correct title and options', () => {
+    const { lastFrame } = render(
+      <ProQuotaDialog
+        currentModel="gemini-2.5-pro"
+        fallbackModel="gemini-2.5-flash"
+        onChoice={() => {}}
+      />,
+    );
+
+    const output = lastFrame();
+    expect(output).toContain('Pro quota limit reached for gemini-2.5-pro.');
+
+    // Check that RadioButtonSelect was called with the correct items
+    expect(RadioButtonSelect).toHaveBeenCalledWith(
+      expect.objectContaining({
+        items: [
+          {
+            label: 'Change auth (executes the /auth command)',
+            value: 'auth',
+          },
+          {
+            label: `Continue with gemini-2.5-flash`,
+            value: 'continue',
+          },
+        ],
+      }),
+      undefined,
+    );
+  });
+
+  it('should call onChoice with "auth" when "Change auth" is selected', () => {
+    const mockOnChoice = vi.fn();
+    render(
+      <ProQuotaDialog
+        currentModel="gemini-2.5-pro"
+        fallbackModel="gemini-2.5-flash"
+        onChoice={mockOnChoice}
+      />,
+    );
+
+    // Get the onSelect function passed to RadioButtonSelect
+    const onSelect = (RadioButtonSelect as Mock).mock.calls[0][0].onSelect;
+
+    // Simulate the selection
+    onSelect('auth');
+
+    expect(mockOnChoice).toHaveBeenCalledWith('auth');
+  });
+
+  it('should call onChoice with "continue" when "Continue with flash" is selected', () => {
+    const mockOnChoice = vi.fn();
+    render(
+      <ProQuotaDialog
+        currentModel="gemini-2.5-pro"
+        fallbackModel="gemini-2.5-flash"
+        onChoice={mockOnChoice}
+      />,
+    );
+
+    // Get the onSelect function passed to RadioButtonSelect
+    const onSelect = (RadioButtonSelect as Mock).mock.calls[0][0].onSelect;
+
+    // Simulate the selection
+    onSelect('continue');
+
+    expect(mockOnChoice).toHaveBeenCalledWith('continue');
+  });
+});
@@ -0,0 +1,52 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import type React from 'react';
+import { Box, Text } from 'ink';
+import { RadioButtonSelect } from './shared/RadioButtonSelect.js';
+import { Colors } from '../colors.js';
+
+interface ProQuotaDialogProps {
+  currentModel: string;
+  fallbackModel: string;
+  onChoice: (choice: 'auth' | 'continue') => void;
+}
+
+export function ProQuotaDialog({
+  currentModel,
+  fallbackModel,
+  onChoice,
+}: ProQuotaDialogProps): React.JSX.Element {
+  const items = [
+    {
+      label: 'Change auth (executes the /auth command)',
+      value: 'auth' as const,
+    },
+    {
+      label: `Continue with ${fallbackModel}`,
+      value: 'continue' as const,
+    },
+  ];
+
+  const handleSelect = (choice: 'auth' | 'continue') => {
+    onChoice(choice);
+  };
+
+  return (
+    <Box borderStyle="round" flexDirection="column" paddingX={1}>
+      <Text bold color={Colors.AccentYellow}>
+        Pro quota limit reached for {currentModel}.
+      </Text>
+      <Box marginTop={1}>
+        <RadioButtonSelect
+          items={items}
+          initialIndex={1}
+          onSelect={handleSelect}
+        />
+      </Box>
+    </Box>
+  );
+}