Add complete handler to Subscription, to handle the case when the ser…

Changelog.md

@@ -2,6 +2,13 @@
 
 ## vNext
 
+### Improvements
+
+- Adds a `onSubscriptionComplete` prop to the `Subscription` component, that
+  can be passed a callback to be called when the subscription observable
+  is completed.  <br/>
+  [@sujeetsr](https://github.com/sujeetsr) in [#2716](https://github.com/apollographql/react-apollo/pull/2716)
+
 ## 2.4.0
 
 ### Bug Fixes

src/Subscriptions.tsx

@@ -29,6 +29,7 @@ export interface SubscriptionProps<TData = any, TVariables = OperationVariables>
   shouldResubscribe?: any;
   client?: ApolloClient<Object>;
   onSubscriptionData?: (options: OnSubscriptionDataOptions<TData>) => any;
+  onSubscriptionComplete?: () => void;
   children?: (result: SubscriptionResult<TData>) => React.ReactNode;
 }
 
@@ -55,6 +56,7 @@ class Subscription<TData = any, TVariables = any> extends React.Component<
     variables: PropTypes.object,
     children: PropTypes.func,
     onSubscriptionData: PropTypes.func,
+    onSubscriptionComplete: PropTypes.func,
     shouldResubscribe: PropTypes.oneOfType([PropTypes.func, PropTypes.bool]),
   };
 
@@ -136,6 +138,7 @@ class Subscription<TData = any, TVariables = any> extends React.Component<
     this.querySubscription = this.queryObservable!.subscribe({
       next: this.updateCurrentData,
       error: this.updateError,
+      complete: this.completeSubscription
     });
   };
 
@@ -165,6 +168,12 @@ class Subscription<TData = any, TVariables = any> extends React.Component<
     });
   };
 
+  private completeSubscription = () => {
+    const { onSubscriptionComplete } = this.props;
+    if (onSubscriptionComplete) onSubscriptionComplete();
+    this.endSubscription();
+  };
+
   private endSubscription = () => {
     if (this.querySubscription) {
       this.querySubscription.unsubscribe();

src/test-links.ts

@@ -124,10 +124,11 @@ export class MockSubscriptionLink extends ApolloLink {
     });
   }
 
-  public simulateResult(result: MockedSubscriptionResult) {
+  public simulateResult(result: MockedSubscriptionResult, complete = false) {
     setTimeout(() => {
       const { observer } = this;
       if (!observer) throw new Error('subscription torn down');
+      if (complete && observer.complete) observer.complete();
       if (result.result && observer.next) observer.next(result.result);
       if (result.error && observer.error) observer.error(result.error);
     }, result.delay || 0);

test/client/Subscription.test.tsx

@@ -130,6 +130,37 @@ it('calls onSubscriptionData if given', done => {
   jest.runTimersToTime(40);
 });
 
+it('should call onSubscriptionComplete if specified', done => {
+  jest.useFakeTimers();
+
+  let count = 0;
+
+  const Component = () => (
+    <Subscription
+      subscription={subscription}
+      onSubscriptionData={() => {
+        count++;
+      }}
+      onSubscriptionComplete={() => {
+        done();
+      }}
+    />
+  );
+
+  wrapper = mount(
+    <ApolloProvider client={client}>
+      <Component />
+    </ApolloProvider>,
+  );
+
+  const interval = setInterval(() => {
+    link.simulateResult(results[count], count === 3);
+    if (count >= 3) clearInterval(interval);
+  }, 10);
+
+  jest.runTimersToTime(40);
+});
+
 it('executes subscription for the variables passed in the props', done => {
   expect.assertions(4);
   const subscriptionWithVariables = gql`