Chore (core): Cleanup force exit handling and test it. Rel to #1237 #…

…1114
SBoudrias · Jun 3, 2023 · 09539fa · 09539fa
1 parent 9652548
commit 09539fa
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 20 deletions.
diff --git a/packages/core/core.test.mts b/packages/core/core.test.mts
@@ -13,38 +13,55 @@ import {
   type KeypressEvent,
 } from './src/index.mjs';
 import stripAnsi from 'strip-ansi';
+import ansiEscapes from 'ansi-escapes';
 
 describe('createPrompt()', () => {
   it('handle async function message', async () => {
-    const viewFunction = vi.fn(() => '');
+    const viewFunction = vi.fn((config, done) => {
+      useEffect(() => {
+        done();
+      }, []);
+
+      return '';
+    });
     const prompt = createPrompt(viewFunction);
     const promise = Promise.resolve('Async message:');
     const renderingDone = render(prompt, { message: () => promise });
 
     // Initially, we leave a few ms for message to resolve
     expect(viewFunction).not.toHaveBeenCalled();
 
-    await renderingDone;
+    const { answer } = await renderingDone;
     expect(viewFunction).toHaveBeenLastCalledWith(
       expect.objectContaining({ message: 'Async message:' }),
       expect.any(Function)
     );
+
+    await answer.catch(() => {});
   });
 
   it('handle deferred message', async () => {
-    const viewFunction = vi.fn(() => '');
+    const viewFunction = vi.fn((config, done) => {
+      useEffect(() => {
+        done();
+      }, []);
+
+      return '';
+    });
     const prompt = createPrompt(viewFunction);
     const promise = Promise.resolve('Async message:');
     const renderingDone = render(prompt, { message: promise });
 
     // Initially, we leave a few ms for message to resolve
     expect(viewFunction).not.toHaveBeenCalled();
 
-    await renderingDone;
+    const { answer } = await renderingDone;
     expect(viewFunction).toHaveBeenLastCalledWith(
       expect.objectContaining({ message: 'Async message:' }),
       expect.any(Function)
     );
+
+    await answer.catch(() => {});
   });
 
   it('onKeypress: allow to implement custom behavior on keypress', async () => {
@@ -252,6 +269,22 @@ describe('Error handling', () => {
       '"useEffect return value must be a cleanup function or nothing."'
     );
   });
+
+  it('cleanup prompt on exit', async () => {
+    const Prompt = () => `Question ${ansiEscapes.cursorHide}`;
+
+    const prompt = createPrompt(Prompt);
+    const { answer, getScreen } = await render(prompt, { message: 'Question' });
+
+    process.emit('SIGINT');
+
+    await expect(answer).rejects.toMatchInlineSnapshot(
+      '[Error: User force closed the prompt with CTRL+C]'
+    );
+
+    const lastScreen = getScreen({ raw: true });
+    expect(lastScreen).toContain(ansiEscapes.cursorShow);
+  });
 });
 
 describe('Separator', () => {

diff --git a/packages/core/src/index.mts b/packages/core/src/index.mts
@@ -139,22 +139,6 @@ export function createPrompt<Value, Config extends AsyncPromptConfig>(
     }) as InquirerReadline;
     const screen = new ScreenManager(rl);
 
-    const onExit = () => {
-      if (context?.clearPromptOnDone) {
-        screen.clean();
-      } else {
-        screen.clearContent();
-      }
-      screen.done();
-
-      process.removeListener('exit', onExit);
-      rl.removeListener('SIGINT', onExit);
-    };
-
-    // Handle cleanup on force exit. Main reason is so we restore the cursor if a prompt hide it.
-    process.on('exit', onExit);
-    rl.on('SIGINT', onExit);
-
     // Set our state before starting the prompt.
     hooks.length = 0;
     sessionRl = rl;
@@ -163,6 +147,29 @@ export function createPrompt<Value, Config extends AsyncPromptConfig>(
     const resolvedConfig = await getPromptConfig(config);
 
     return new Promise((resolve, reject) => {
+      const onExit = () => {
+        if (context?.clearPromptOnDone) {
+          screen.clean();
+        } else {
+          screen.clearContent();
+        }
+        screen.done();
+
+        process.removeListener('SIGINT', onForceExit);
+      };
+
+      let shouldHandleExit = true;
+      const onForceExit = () => {
+        if (shouldHandleExit) {
+          shouldHandleExit = false;
+          onExit();
+          reject(new Error('User force closed the prompt with CTRL+C'));
+        }
+      };
+
+      // Handle cleanup on force exit. Main reason is so we restore the cursor if a prompt hide it.
+      process.on('SIGINT', onForceExit);
+
       const done = (value: Value) => {
         // Delay execution to let time to the hookCleanup functions to registers.
         setImmediate(() => {