From ee15dfc40d01c6d8e01438380b115025e4ddc427 Mon Sep 17 00:00:00 2001
From: wolvenstone <stefan.wolfsteiner@gmail.com>
Date: Sat, 8 Jul 2023 14:48:33 +0200
Subject: [PATCH] Let users cancel abortable functions

---
 src/index.ts          | 29 +++++++++++++++++------
 test/useAsync.test.ts | 55 +++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 75 insertions(+), 9 deletions(-)

diff --git a/src/index.ts b/src/index.ts
index 05c52c9..c61a66f 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -224,6 +224,13 @@ export type UseAsyncReturn<
   currentParams: Args | null;
 };
 
+export type UseAsyncAbortableReturn<
+  R = UnknownResult,
+  Args extends any[] = UnknownArgs
+> = UseAsyncReturn<R, Args> & {
+  cancel: () => void;
+};
+
 // Relaxed interface which accept both async and sync functions
 // Accepting sync function is convenient for useAsyncCallback
 const useAsyncInternal = <R = UnknownResult, Args extends any[] = UnknownArgs>(
@@ -327,9 +334,9 @@ export function useAsync<R = UnknownResult, Args extends any[] = UnknownArgs>(
   return useAsyncInternal(asyncFunction, params, options);
 }
 
-type AddArg<H, T extends any[]> = ((h: H, ...t: T) => void) extends ((
+type AddArg<H, T extends any[]> = ((h: H, ...t: T) => void) extends (
   ...r: infer R
-) => void)
+) => void
   ? R
   : never;
 
@@ -340,17 +347,22 @@ export const useAsyncAbortable = <
   asyncFunction: (...args: AddArg<AbortSignal, Args>) => Promise<R>,
   params: Args,
   options?: UseAsyncOptions<R>
-): UseAsyncReturn<R, Args> => {
+): UseAsyncAbortableReturn<R, Args> => {
   const abortControllerRef = useRef<AbortController>();
 
+  const abortHandler = () => {
+    if (abortControllerRef.current) {
+      abortControllerRef.current.abort();
+    }
+  };
+
   // Wrap the original async function and enhance it with abortion login
   const asyncFunctionWrapper: (...args: Args) => Promise<R> = async (
     ...p: Args
   ) => {
     // Cancel previous async call
-    if (abortControllerRef.current) {
-      abortControllerRef.current.abort();
-    }
+    abortHandler();
+
     // Create/store new abort controller for next async call
     const abortController = new AbortController();
     abortControllerRef.current = abortController;
@@ -367,7 +379,10 @@ export const useAsyncAbortable = <
     }
   };
 
-  return useAsync(asyncFunctionWrapper, params, options);
+  return {
+    ...useAsync(asyncFunctionWrapper, params, options),
+    cancel: () => abortHandler(),
+  };
 };
 
 // keep compat with TS < 3.5
diff --git a/test/useAsync.test.ts b/test/useAsync.test.ts
index 7e163c3..038e3f3 100644
--- a/test/useAsync.test.ts
+++ b/test/useAsync.test.ts
@@ -1,7 +1,25 @@
-import { useAsync } from '../src';
+import { useAsync, useAsyncAbortable } from '../src';
 import { renderHook } from '@testing-library/react-hooks';
 
-const sleep = (ms: number) => new Promise(resolve => setTimeout(resolve, ms));
+const sleep = (ms: number, signal?: AbortSignal) =>
+  new Promise((resolve, reject) => {
+    if (signal && signal.aborted) reject();
+
+    const timeout = setTimeout(() => {
+      resolve(undefined);
+
+      if (signal) signal.removeEventListener('abort', abort);
+    }, ms);
+
+    const abort = () => {
+      if (signal) {
+        clearTimeout(timeout);
+        reject();
+      }
+    };
+
+    if (signal) signal.addEventListener('abort', abort);
+  });
 
 interface StarwarsHero {
   name: string;
@@ -269,4 +287,37 @@ describe('useAync', () => {
     expect(onSuccess).not.toHaveBeenCalled();
     expect(onError).toHaveBeenCalled();
   });
+
+  it('should handle cancel', async () => {
+    const onSuccess = jest.fn();
+    const onError = jest.fn();
+    const onAbort = jest.fn();
+
+    const { result, waitForNextUpdate } = renderHook(() =>
+      useAsyncAbortable(
+        async (signal: AbortSignal) => {
+          signal.addEventListener('abort', onAbort);
+          await sleep(10000, signal);
+          return fakeResults;
+        },
+        [],
+        {
+          onSuccess: () => onSuccess(),
+          onError: () => onError(),
+        }
+      )
+    );
+
+    await sleep(10);
+    expect(result.current.loading).toBe(true);
+    result.current.cancel();
+    await waitForNextUpdate();
+
+    expect(onAbort).toHaveBeenCalled();
+    expect(result.current.error).toBeUndefined();
+    expect(result.current.loading).toBe(false);
+    expect(result.current.result).toBeUndefined();
+    expect(onSuccess).not.toHaveBeenCalled();
+    expect(onError).toHaveBeenCalled();
+  });
 });