Refactoring packages/cli/src/ui tests (#12482)

Co-authored-by: riddhi <duttariddhi@google.com>
This commit is contained in:
Riddhi Dutta
2025-11-03 23:40:57 +05:30
committed by GitHub
parent 93f14ce626
commit 19ea68b838
5 changed files with 810 additions and 1093 deletions
+145 -223
View File
@@ -201,6 +201,28 @@ describe('useReactToolScheduler', () => {
| ((outcome: ToolConfirmationOutcome) => void | Promise<void>)
| undefined;
const advanceAndSettle = async () => {
await act(async () => {
await vi.advanceTimersByTimeAsync(0);
});
};
const scheduleAndWaitForExecution = async (
schedule: (
req: ToolCallRequestInfo | ToolCallRequestInfo[],
signal: AbortSignal,
) => void,
request: ToolCallRequestInfo | ToolCallRequestInfo[],
) => {
act(() => {
schedule(request, new AbortController().signal);
});
await advanceAndSettle();
await advanceAndSettle();
await advanceAndSettle();
};
beforeEach(() => {
onComplete = vi.fn();
capturedOnConfirmForTest = undefined;
@@ -259,7 +281,6 @@ describe('useReactToolScheduler', () => {
(mockTool.shouldConfirmExecute as Mock).mockResolvedValue(null);
const { result } = renderScheduler();
const schedule = result.current[1];
const request: ToolCallRequestInfo = {
callId: 'call1',
name: 'mockTool',
@@ -271,42 +292,19 @@ describe('useReactToolScheduler', () => {
completedToolCalls = calls;
});
act(() => {
schedule(request, new AbortController().signal);
});
await act(async () => {
await vi.advanceTimersByTimeAsync(0);
});
await act(async () => {
await vi.advanceTimersByTimeAsync(0);
});
await act(async () => {
await vi.advanceTimersByTimeAsync(0);
});
await scheduleAndWaitForExecution(result.current[1], request);
expect(mockTool.execute).toHaveBeenCalledWith(request.args);
expect(onComplete).toHaveBeenCalledWith([
expect.objectContaining({
status: 'success',
request,
response: expect.objectContaining({
resultDisplay: 'Formatted tool output',
responseParts: [
{
functionResponse: {
id: 'call1',
name: 'mockTool',
response: { output: 'Tool output' },
},
},
],
}),
}),
]);
expect(completedToolCalls).toHaveLength(1);
expect(completedToolCalls[0].status).toBe('success');
expect(completedToolCalls[0].request).toBe(request);
if (
completedToolCalls[0].status === 'success' ||
completedToolCalls[0].status === 'error'
) {
expect(completedToolCalls[0].response).toMatchSnapshot();
}
});
it('should clear previous tool calls when scheduling new ones', async () => {
@@ -412,113 +410,83 @@ describe('useReactToolScheduler', () => {
});
});
it('should handle tool not found', async () => {
mockToolRegistry.getTool.mockReturnValue(undefined);
const { result } = renderScheduler();
const schedule = result.current[1];
const request: ToolCallRequestInfo = {
callId: 'call1',
name: 'nonexistentTool',
args: {},
} as any;
it.each([
{
desc: 'tool not found',
setup: () => {
mockToolRegistry.getTool.mockReturnValue(undefined);
},
request: {
callId: 'call1',
name: 'nonexistentTool',
args: {},
} as any,
expectedErrorContains: [
'Tool "nonexistentTool" not found in registry',
'Did you mean one of:',
],
},
{
desc: 'error during shouldConfirmExecute',
setup: () => {
mockToolRegistry.getTool.mockReturnValue(mockTool);
const confirmError = new Error('Confirmation check failed');
(mockTool.shouldConfirmExecute as Mock).mockRejectedValue(confirmError);
},
request: {
callId: 'call1',
name: 'mockTool',
args: {},
} as any,
expectedError: new Error('Confirmation check failed'),
},
{
desc: 'error during execute',
setup: () => {
mockToolRegistry.getTool.mockReturnValue(mockTool);
(mockTool.shouldConfirmExecute as Mock).mockResolvedValue(null);
const execError = new Error('Execution failed');
(mockTool.execute as Mock).mockRejectedValue(execError);
},
request: {
callId: 'call1',
name: 'mockTool',
args: {},
} as any,
expectedError: new Error('Execution failed'),
},
])(
'should handle $desc',
async ({ setup, request, expectedErrorContains, expectedError }) => {
setup();
const { result } = renderScheduler();
let completedToolCalls: ToolCall[] = [];
onComplete.mockImplementation((calls) => {
completedToolCalls = calls;
});
let completedToolCalls: ToolCall[] = [];
onComplete.mockImplementation((calls) => {
completedToolCalls = calls;
});
act(() => {
schedule(request, new AbortController().signal);
});
await act(async () => {
await vi.advanceTimersByTimeAsync(0);
});
await act(async () => {
await vi.advanceTimersByTimeAsync(0);
});
await scheduleAndWaitForExecution(result.current[1], request);
expect(completedToolCalls).toHaveLength(1);
expect(completedToolCalls[0].status).toBe('error');
expect(completedToolCalls[0].request).toBe(request);
expect((completedToolCalls[0] as any).response.error.message).toContain(
'Tool "nonexistentTool" not found in registry',
);
expect((completedToolCalls[0] as any).response.error.message).toContain(
'Did you mean one of:',
);
});
expect(completedToolCalls).toHaveLength(1);
expect(completedToolCalls[0].status).toBe('error');
expect(completedToolCalls[0].request).toBe(request);
it('should handle error during shouldConfirmExecute', async () => {
mockToolRegistry.getTool.mockReturnValue(mockTool);
const confirmError = new Error('Confirmation check failed');
(mockTool.shouldConfirmExecute as Mock).mockRejectedValue(confirmError);
if (expectedErrorContains) {
expectedErrorContains.forEach((errorText) => {
expect(
(completedToolCalls[0] as any).response.error.message,
).toContain(errorText);
});
}
const { result } = renderScheduler();
const schedule = result.current[1];
const request: ToolCallRequestInfo = {
callId: 'call1',
name: 'mockTool',
args: {},
} as any;
let completedToolCalls: ToolCall[] = [];
onComplete.mockImplementation((calls) => {
completedToolCalls = calls;
});
act(() => {
schedule(request, new AbortController().signal);
});
await act(async () => {
await vi.advanceTimersByTimeAsync(0);
});
await act(async () => {
await vi.advanceTimersByTimeAsync(0);
});
expect(completedToolCalls).toHaveLength(1);
expect(completedToolCalls[0].status).toBe('error');
expect(completedToolCalls[0].request).toBe(request);
expect((completedToolCalls[0] as any).response.error).toBe(confirmError);
});
it('should handle error during execute', async () => {
mockToolRegistry.getTool.mockReturnValue(mockTool);
(mockTool.shouldConfirmExecute as Mock).mockResolvedValue(null);
const execError = new Error('Execution failed');
(mockTool.execute as Mock).mockRejectedValue(execError);
const { result } = renderScheduler();
const schedule = result.current[1];
const request: ToolCallRequestInfo = {
callId: 'call1',
name: 'mockTool',
args: {},
} as any;
let completedToolCalls: ToolCall[] = [];
onComplete.mockImplementation((calls) => {
completedToolCalls = calls;
});
act(() => {
schedule(request, new AbortController().signal);
});
await act(async () => {
await vi.advanceTimersByTimeAsync(0);
});
await act(async () => {
await vi.advanceTimersByTimeAsync(0);
});
await act(async () => {
await vi.advanceTimersByTimeAsync(0);
});
expect(completedToolCalls).toHaveLength(1);
expect(completedToolCalls[0].status).toBe('error');
expect(completedToolCalls[0].request).toBe(request);
expect((completedToolCalls[0] as any).response.error).toBe(execError);
});
if (expectedError) {
expect((completedToolCalls[0] as any).response.error.message).toBe(
expectedError.message,
);
}
},
);
it('should handle tool requiring confirmation - approved', async () => {
mockToolRegistry.getTool.mockReturnValue(mockToolRequiresConfirmation);
@@ -539,9 +507,7 @@ describe('useReactToolScheduler', () => {
act(() => {
schedule(request, new AbortController().signal);
});
await act(async () => {
await vi.advanceTimersByTimeAsync(0);
});
await advanceAndSettle();
const waitingCall = result.current[0][0] as any;
expect(waitingCall.status).toBe('awaiting_approval');
@@ -552,36 +518,24 @@ describe('useReactToolScheduler', () => {
await capturedOnConfirmForTest?.(ToolConfirmationOutcome.ProceedOnce);
});
await act(async () => {
await vi.advanceTimersByTimeAsync(0);
});
await act(async () => {
await vi.advanceTimersByTimeAsync(0);
});
await act(async () => {
await vi.advanceTimersByTimeAsync(0);
});
await advanceAndSettle();
await advanceAndSettle();
await advanceAndSettle();
expect(mockOnUserConfirmForToolConfirmation).toHaveBeenCalledWith(
ToolConfirmationOutcome.ProceedOnce,
);
expect(mockToolRequiresConfirmation.execute).toHaveBeenCalled();
expect(onComplete).toHaveBeenCalledWith([
expect.objectContaining({
status: 'success',
request,
response: expect.objectContaining({
resultDisplay: 'Confirmed display',
responseParts: expect.arrayContaining([
expect.objectContaining({
functionResponse: expect.objectContaining({
response: { output: expectedOutput },
}),
}),
]),
}),
}),
]);
const completedCalls = onComplete.mock.calls[0][0] as ToolCall[];
expect(completedCalls[0].status).toBe('success');
expect(completedCalls[0].request).toBe(request);
if (
completedCalls[0].status === 'success' ||
completedCalls[0].status === 'error'
) {
expect(completedCalls[0].response).toMatchSnapshot();
}
});
it('should handle tool requiring confirmation - cancelled by user', async () => {
@@ -597,9 +551,7 @@ describe('useReactToolScheduler', () => {
act(() => {
schedule(request, new AbortController().signal);
});
await act(async () => {
await vi.advanceTimersByTimeAsync(0);
});
await advanceAndSettle();
const waitingCall = result.current[0][0] as any;
expect(waitingCall.status).toBe('awaiting_approval');
@@ -609,34 +561,23 @@ describe('useReactToolScheduler', () => {
await act(async () => {
await capturedOnConfirmForTest?.(ToolConfirmationOutcome.Cancel);
});
await act(async () => {
await vi.advanceTimersByTimeAsync(0);
});
await act(async () => {
await vi.advanceTimersByTimeAsync(0);
});
await advanceAndSettle();
await advanceAndSettle();
expect(mockOnUserConfirmForToolConfirmation).toHaveBeenCalledWith(
ToolConfirmationOutcome.Cancel,
);
expect(onComplete).toHaveBeenCalledWith([
expect.objectContaining({
status: 'cancelled',
request,
response: expect.objectContaining({
responseParts: expect.arrayContaining([
expect.objectContaining({
functionResponse: expect.objectContaining({
response: expect.objectContaining({
error:
'[Operation Cancelled] Reason: User cancelled the operation.',
}),
}),
}),
]),
}),
}),
]);
const completedCalls = onComplete.mock.calls[0][0] as ToolCall[];
expect(completedCalls[0].status).toBe('cancelled');
expect(completedCalls[0].request).toBe(request);
if (
completedCalls[0].status === 'success' ||
completedCalls[0].status === 'error' ||
completedCalls[0].status === 'cancelled'
) {
expect(completedCalls[0].response).toMatchSnapshot();
}
});
it('should handle live output updates', async () => {
@@ -662,7 +603,6 @@ describe('useReactToolScheduler', () => {
);
const { result } = renderScheduler();
const schedule = result.current[1];
const request: ToolCallRequestInfo = {
callId: 'liveCall',
name: 'mockToolWithLiveOutput',
@@ -670,11 +610,9 @@ describe('useReactToolScheduler', () => {
} as any;
act(() => {
schedule(request, new AbortController().signal);
});
await act(async () => {
await vi.advanceTimersByTimeAsync(0);
result.current[1](request, new AbortController().signal);
});
await advanceAndSettle();
expect(liveUpdateFn).toBeDefined();
expect(result.current[0][0].status).toBe('executing');
@@ -682,16 +620,12 @@ describe('useReactToolScheduler', () => {
await act(async () => {
liveUpdateFn?.('Live output 1');
});
await act(async () => {
await vi.advanceTimersByTimeAsync(0);
});
await advanceAndSettle();
await act(async () => {
liveUpdateFn?.('Live output 2');
});
await act(async () => {
await vi.advanceTimersByTimeAsync(0);
});
await advanceAndSettle();
act(() => {
resolveExecutePromise({
@@ -699,29 +633,18 @@ describe('useReactToolScheduler', () => {
returnDisplay: 'Final display',
} as ToolResult);
});
await act(async () => {
await vi.advanceTimersByTimeAsync(0);
});
await act(async () => {
await vi.advanceTimersByTimeAsync(0);
});
await advanceAndSettle();
await advanceAndSettle();
expect(onComplete).toHaveBeenCalledWith([
expect.objectContaining({
status: 'success',
request,
response: expect.objectContaining({
resultDisplay: 'Final display',
responseParts: expect.arrayContaining([
expect.objectContaining({
functionResponse: expect.objectContaining({
response: { output: 'Final output' },
}),
}),
]),
}),
}),
]);
const completedCalls = onComplete.mock.calls[0][0] as ToolCall[];
expect(completedCalls[0].status).toBe('success');
expect(completedCalls[0].request).toBe(request);
if (
completedCalls[0].status === 'success' ||
completedCalls[0].status === 'error'
) {
expect(completedCalls[0].response).toMatchSnapshot();
}
expect(result.current[0]).toEqual([]);
});
@@ -874,7 +797,6 @@ describe('useReactToolScheduler', () => {
response: expect.objectContaining({ resultDisplay: 'done display' }),
}),
]);
// Wait for request2 to complete
await act(async () => {
await vi.advanceTimersByTimeAsync(50);
await vi.advanceTimersByTimeAsync(0);