From 53dc100c35daf8fe6f033c81e42164d667f2f85f Mon Sep 17 00:00:00 2001
From: Sairaj Chouhan <sairaj2119@gmail.com>
Date: Thu, 27 Jan 2022 20:30:30 +0530
Subject: [PATCH] feat: shouldErrorRetry accepts a function (#1816)

---
 src/types.ts                |  2 +-
 src/use-swr.ts              |  7 ++++-
 test/use-swr-error.test.tsx | 56 +++++++++++++++++++++++++++++++++++++
 3 files changed, 63 insertions(+), 2 deletions(-)

diff --git a/src/types.ts b/src/types.ts
index 32f0e81ae..058fdf35d 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -42,7 +42,7 @@ export interface PublicConfiguration<
   revalidateOnReconnect: boolean
   revalidateOnMount?: boolean
   revalidateIfStale: boolean
-  shouldRetryOnError: boolean
+  shouldRetryOnError: boolean | ((err: Error) => boolean)
   suspense?: boolean
   fallbackData?: Data
   fetcher?: Fn
diff --git a/src/use-swr.ts b/src/use-swr.ts
index 3dc30fc77..1b28b0467 100644
--- a/src/use-swr.ts
+++ b/src/use-swr.ts
@@ -304,7 +304,12 @@ export const useSWRHandler = <Data = any, Error = any>(
           // deduped ones.
           if (shouldStartNewRequest && isCurrentKeyMounted()) {
             getConfig().onError(err, key, config)
-            if (config.shouldRetryOnError) {
+            if (
+              (typeof config.shouldRetryOnError === 'boolean' &&
+                config.shouldRetryOnError) ||
+              (isFunction(config.shouldRetryOnError) &&
+                config.shouldRetryOnError(err as Error))
+            ) {
               // When retrying, dedupe is always enabled
               if (isActive()) {
                 // If it's active, stop. It will auto revalidate when refocusing
diff --git a/test/use-swr-error.test.tsx b/test/use-swr-error.test.tsx
index 249ada79c..ccdaf35f6 100644
--- a/test/use-swr-error.test.tsx
+++ b/test/use-swr-error.test.tsx
@@ -142,6 +142,62 @@ describe('useSWR - error', () => {
     screen.getByText('error: 0')
   })
 
+  it('should not retry when shouldRetryOnError function returns false', async () => {
+    const key = createKey()
+    let count = 0
+    function Page() {
+      const { data, error } = useSWR(
+        key,
+        () => createResponse(new Error('error: ' + count++), { delay: 100 }),
+        {
+          onErrorRetry: (_, __, ___, revalidate, revalidateOpts) => {
+            revalidate(revalidateOpts)
+          },
+          dedupingInterval: 0,
+          shouldRetryOnError: () => false
+        }
+      )
+      if (error) return <div>{error.message}</div>
+      return <div>hello, {data}</div>
+    }
+    renderWithConfig(<Page />)
+    screen.getByText('hello,')
+
+    // mount
+    await screen.findByText('error: 0')
+
+    await act(() => sleep(150))
+    screen.getByText('error: 0')
+  })
+
+  it('should retry when shouldRetryOnError function returns true', async () => {
+    const key = createKey()
+    let count = 0
+    function Page() {
+      const { data, error } = useSWR(
+        key,
+        () => createResponse(new Error('error: ' + count++), { delay: 100 }),
+        {
+          onErrorRetry: (_, __, ___, revalidate, revalidateOpts) => {
+            revalidate(revalidateOpts)
+          },
+          dedupingInterval: 0,
+          shouldRetryOnError: () => true
+        }
+      )
+      if (error) return <div>{error.message}</div>
+      return <div>hello, {data}</div>
+    }
+    renderWithConfig(<Page />)
+    screen.getByText('hello,')
+
+    // mount
+    await screen.findByText('error: 0')
+
+    await act(() => sleep(150))
+    screen.getByText('error: 1')
+  })
+
   it('should trigger the onLoadingSlow and onSuccess event', async () => {
     const key = createKey()
     let loadingSlow = null,