From ad850570cb04b3c8434c6f528f9de756f38aae60 Mon Sep 17 00:00:00 2001
From: Jeroen Visser <jeroenvisser101@gmail.com>
Date: Tue, 8 Mar 2022 00:56:29 +0100
Subject: [PATCH] Add onError option to useSubscription hook

---
 .../hooks/__tests__/useSubscription.test.tsx  | 53 ++++++++++++++++++-
 src/react/hooks/useSubscription.ts            |  1 +
 src/react/types/types.ts                      |  1 +
 3 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/src/react/hooks/__tests__/useSubscription.test.tsx b/src/react/hooks/__tests__/useSubscription.test.tsx
index caeb9949f06..1c8f26da2f3 100644
--- a/src/react/hooks/__tests__/useSubscription.test.tsx
+++ b/src/react/hooks/__tests__/useSubscription.test.tsx
@@ -2,7 +2,7 @@ import React from 'react';
 import { renderHook } from '@testing-library/react-hooks';
 import gql from 'graphql-tag';
 
-import { ApolloClient, ApolloLink, concat } from '../../../core';
+import { ApolloClient, ApolloError, ApolloLink, concat } from '../../../core';
 import { InMemoryCache as Cache } from '../../../cache';
 import { ApolloProvider } from '../../context';
 import { MockSubscriptionLink } from '../../../testing';
@@ -61,6 +61,57 @@ describe('useSubscription Hook', () => {
     expect(result.current.data).toEqual(results[3].result.data);
   });
 
+  it('should call onError after error results', async () => {
+    const subscription = gql`
+      subscription {
+        car {
+          make
+        }
+      }
+    `;
+
+    const results = ['Audi', 'BMW', 'Mercedes', 'Hyundai'].map(make => ({
+      result: { data: { car: { make } } }
+    }));
+
+    const errorResult = {
+      error: new ApolloError({ errorMessage: "test" }),
+      result: { data: { car: { make: null } } },
+    };
+
+    const link = new MockSubscriptionLink();
+    const client = new ApolloClient({
+      link,
+      cache: new Cache({ addTypename: false })
+    });
+
+
+    const onError = jest.fn();
+    const { result, waitForNextUpdate } = renderHook(
+      () => useSubscription(subscription, { onError }),
+      {
+        wrapper: ({ children }) => (
+          <ApolloProvider client={client}>
+            {children}
+          </ApolloProvider>
+        ),
+      },
+    );
+
+    expect(result.current.loading).toBe(true);
+    expect(result.current.error).toBe(undefined);
+    expect(result.current.data).toBe(undefined);
+    setTimeout(() => link.simulateResult(results[0]));
+    await waitForNextUpdate();
+    expect(result.current.loading).toBe(false);
+    expect(result.current.data).toEqual(results[0].result.data);
+    setTimeout(() => link.simulateResult(errorResult));
+    await waitForNextUpdate();
+    setTimeout(() => {
+      expect(onError).toHaveBeenCalledTimes(1);
+    });
+  });
+
   it('should cleanup after the subscription component has been unmounted', async () => {
     const subscription = gql`
       subscription {
diff --git a/src/react/hooks/useSubscription.ts b/src/react/hooks/useSubscription.ts
index f1335048b53..a11b9944a0f 100644
--- a/src/react/hooks/useSubscription.ts
+++ b/src/react/hooks/useSubscription.ts
@@ -110,6 +110,7 @@ export function useSubscription<TData = any, TVariables = OperationVariables>(
           error,
           variables: options?.variables,
         });
+        ref.current.options?.onError?.(error);
       },
       complete() {
         ref.current.options?.onSubscriptionComplete?.();
diff --git a/src/react/types/types.ts b/src/react/types/types.ts
index f74817d8117..60ead990d02 100644
--- a/src/react/types/types.ts
+++ b/src/react/types/types.ts
@@ -211,6 +211,7 @@ export interface BaseSubscriptionOptions<
   skip?: boolean;
   context?: DefaultContext;
   onSubscriptionData?: (options: OnSubscriptionDataOptions<TData>) => any;
+  onError?: (error: ApolloError) => void;
   onSubscriptionComplete?: () => void;
 }