diff --git a/packages/cli-kit/src/private/node/ui/components/Tasks.test.tsx b/packages/cli-kit/src/private/node/ui/components/Tasks.test.tsx index 7fb3986930..6d8f254a3e 100644 --- a/packages/cli-kit/src/private/node/ui/components/Tasks.test.tsx +++ b/packages/cli-kit/src/private/node/ui/components/Tasks.test.tsx @@ -1,5 +1,5 @@ import {Task, Tasks} from './Tasks.js' -import {render} from '../../testing/ui.js' +import {getLastFrameAfterUnmount, render, waitForContent} from '../../testing/ui.js' import {TokenizedString} from '../../../../public/node/output.js' import {AbortController} from '../../../../public/node/abort.js' import {Stdout} from '../../ui.js' @@ -29,33 +29,43 @@ beforeEach(() => { describe('Tasks', () => { test('shows nothing at the end in case of success', async () => { // Given - const firstTaskFunction = vi.fn(async () => {}) - const secondTaskFunction = vi.fn(async () => {}) + let resolveTask!: () => void + const firstTaskFunction = vi.fn(async () => { + await new Promise((resolve) => { + resolveTask = resolve + }) + }) const firstTask = { title: 'task 1', task: firstTaskFunction, } - const secondTask = { - title: 'task 2', - task: secondTaskFunction, - } // When + const renderInstance = render() + await waitForContent(renderInstance, 'task 1') - const renderInstance = render() + resolveTask() await renderInstance.waitUntilExit() + + // Then + expect(firstTaskFunction).toHaveBeenCalledTimes(1) + expect(getLastFrameAfterUnmount(renderInstance)).toBe('') }) test('stops at the task that throws error', async () => { // Given const abortController = new AbortController() const secondTaskFunction = vi.fn(async () => {}) + const error = new Error('something went wrong') + let rejectTask!: (error: Error) => void const firstTask: Task = { title: 'task 1', task: async () => { - throw new Error('something went wrong') + await new Promise((_resolve, reject) => { + rejectTask = reject + }) }, } @@ -68,10 +78,15 @@ describe('Tasks', () => { const renderInstance = render( , ) + await waitForContent(renderInstance, 'task 1') + + const exitPromise = renderInstance.waitUntilExit() + rejectTask(error) // Then - await expect(renderInstance.waitUntilExit()).rejects.toThrowError('something went wrong') + await expect(exitPromise).rejects.toThrowError('something went wrong') expect(secondTaskFunction).toHaveBeenCalledTimes(0) + expect(getLastFrameAfterUnmount(renderInstance)).toBe('') }) test('it supports subtasks', async () => { diff --git a/packages/cli-kit/src/private/node/ui/hooks/use-async-and-unmount.ts b/packages/cli-kit/src/private/node/ui/hooks/use-async-and-unmount.ts index e9a3889047..8588bf8a00 100644 --- a/packages/cli-kit/src/private/node/ui/hooks/use-async-and-unmount.ts +++ b/packages/cli-kit/src/private/node/ui/hooks/use-async-and-unmount.ts @@ -12,17 +12,21 @@ export default function useAsyncAndUnmount( ) { const {exit: unmountInk} = useApp() + const scheduleUnmount = (error?: Error) => { + // Defer unmounting to the next setImmediate so React 19 can flush + // batched state updates before the tree is torn down. + setImmediate(() => unmountInk(error)) + } + useEffect(() => { asyncFunction() .then(() => { onFulfilled() - // Defer unmount so React 19 can flush batched state updates - // before the component tree is torn down. - setImmediate(() => unmountInk()) + scheduleUnmount() }) .catch((error) => { onRejected(error) - setImmediate(() => unmountInk(error)) + scheduleUnmount(error) }) }, []) }