AAE-51358 Cleaning up websocket service from old subscription protocol (#12212)

* AAE-51358 Cleaning up websocket service from old subscription protocol

* AAE-51358 Fixing retry on error bug

lib/process-services-cloud/src/lib/services/web-socket.service.spec.ts

@@ -16,18 +16,39 @@
  */
 
 import { TestBed } from '@angular/core/testing';
+import { Injectable } from '@angular/core';
 import { Apollo, gql } from 'apollo-angular';
 import { lastValueFrom, of, Subject } from 'rxjs';
 import { WebSocketService } from './web-socket.service';
-import { SubscriptionOptions } from '@apollo/client/core';
+import { ApolloLink, execute, FetchResult, Observable as ApolloObservable, SubscriptionOptions } from '@apollo/client/core';
 import { provideHttpClientTesting } from '@angular/common/http/testing';
 import { AuthenticationService, AppConfigService } from '@alfresco/adf-core';
+import { Client, ClientOptions, Sink, SubscribePayload } from 'graphql-ws';
+import { HttpLink } from 'apollo-angular/http';
+
+@Injectable()
+class TestWebSocketService extends WebSocketService {
+    public capturedOnError: (() => void) | undefined;
+
+    protected override createWsClient(clientOptions: ClientOptions): Client {
+        this.capturedOnError = clientOptions.on?.error as (() => void) | undefined;
+
+        return {
+            on: () => () => undefined,
+            subscribe: (_payload: SubscribePayload, _sink: Sink) => () => undefined,
+            async *iterate() {},
+            terminate: () => undefined,
+            dispose: () => undefined
+        };
+    }
+}
 
 describe('WebSocketService', () => {
-    let service: WebSocketService;
+    let service: TestWebSocketService;
     const onLogoutSubject: Subject<void> = new Subject<void>();
 
-    const apolloMock = jasmine.createSpyObj('Apollo', ['use', 'createNamed']);
+    const apolloMock = jasmine.createSpyObj('Apollo', ['use', 'createNamed', 'removeClient']);
+    const httpLinkMock = jasmine.createSpyObj('HttpLink', ['create']);
 
     beforeEach(() => {
         TestBed.configureTestingModule({
@@ -37,6 +58,14 @@ describe('WebSocketService', () => {
                     provide: Apollo,
                     useValue: apolloMock
                 },
+                {
+                    provide: WebSocketService,
+                    useClass: TestWebSocketService
+                },
+                {
+                    provide: HttpLink,
+                    useValue: httpLinkMock
+                },
                 {
                     provide: AppConfigService,
                     useValue: {
@@ -52,13 +81,15 @@ describe('WebSocketService', () => {
                 }
             ]
         });
-        service = TestBed.inject(WebSocketService);
+        service = TestBed.inject(WebSocketService) as TestWebSocketService;
         apolloMock.use.and.returnValues(undefined, { subscribe: () => of({}) });
     });
 
     afterEach(() => {
         apolloMock.use.calls.reset();
         apolloMock.createNamed.calls.reset();
+        apolloMock.removeClient.calls.reset();
+        httpLinkMock.create.calls.reset();
     });
 
     it('should not create a new Apollo client if it is already in use', async () => {
@@ -95,7 +126,7 @@ describe('WebSocketService', () => {
         const apolloClientName = 'testClient';
         const subscriptionOptions: SubscriptionOptions = { query: gql(`subscription {testQuery}`) };
         const wsOptions = { apolloClientName, wsUrl: 'testUrl', subscriptionOptions };
-        apolloMock.createNamed.and.callFake((_, options) => {
+        apolloMock.createNamed.and.callFake((_: any, options: { headers: {} }) => {
             headers = options.headers;
         });
 
@@ -105,4 +136,74 @@ describe('WebSocketService', () => {
         expect(apolloMock.createNamed).toHaveBeenCalled();
         expect(headers).toEqual(expectedHeaders);
     });
+
+    it('should recreate the subscription client when the websocket connection errors', async () => {
+        const apolloClientName = 'testClient';
+        const subscriptionOptions: SubscriptionOptions = { query: gql(`subscription {testQuery}`) };
+        const wsOptions = { apolloClientName, wsUrl: 'testUrl', subscriptionOptions };
+
+        await lastValueFrom(service.getSubscription(wsOptions));
+
+        expect(apolloMock.createNamed).toHaveBeenCalledTimes(1);
+        expect(apolloMock.removeClient).not.toHaveBeenCalled();
+
+        if (!service.capturedOnError) {
+            fail('Expected websocket error handler to be registered');
+            return;
+        }
+
+        service.capturedOnError();
+
+        expect(apolloMock.removeClient).toHaveBeenCalledWith(apolloClientName);
+        expect(apolloMock.createNamed).toHaveBeenCalledTimes(2);
+        expect(apolloMock.createNamed).toHaveBeenCalledWith(apolloClientName, jasmine.any(Object));
+    });
+
+    it('should retry the operation when a GraphQL error is unauthenticated', async () => {
+        const apolloClientName = 'testClient';
+        const subscriptionOptions: SubscriptionOptions = { query: gql(`subscription {testQuery}`) };
+        const wsOptions = { apolloClientName, wsUrl: 'testUrl', httpUrl: 'testHttpUrl', subscriptionOptions };
+        const expectedResult: FetchResult = { data: { retried: true } };
+        let createdLink: ApolloLink | undefined;
+        let requestCount = 0;
+
+        httpLinkMock.create.and.returnValue(
+            new ApolloLink(
+                () =>
+                    new ApolloObservable<FetchResult>((observer) => {
+                        requestCount++;
+
+                        if (requestCount === 1) {
+                            observer.next({
+                                errors: [{ message: 'Unauthorized', extensions: { code: 'UNAUTHENTICATED' } }]
+                            });
+                        } else {
+                            observer.next(expectedResult);
+                        }
+
+                        observer.complete();
+                    })
+            )
+        );
+        apolloMock.createNamed.and.callFake((_clientName: any, options: { link: ApolloLink | undefined }) => {
+            createdLink = options.link;
+        });
+
+        await lastValueFrom(service.getSubscription(wsOptions));
+
+        if (!createdLink) {
+            fail('Expected Apollo link to be created');
+            return;
+        }
+
+        const result = await new Promise<FetchResult>((resolve, reject) => {
+            execute(createdLink!, { query: gql(`query { testQuery }`) }).subscribe({
+                next: resolve,
+                error: reject
+            });
+        });
+
+        expect(requestCount).toBe(2);
+        expect(result).toEqual(expectedResult);
+    });
 });

lib/process-services-cloud/src/lib/services/web-socket.service.ts

@@ -15,10 +15,9 @@
  * limitations under the License.
  */
 
-import { createClient } from 'graphql-ws';
+import { Client, ClientOptions, createClient } from 'graphql-ws';
 import { inject, Injectable } from '@angular/core';
 import { GraphQLWsLink } from '@apollo/client/link/subscriptions';
-import { WebSocketLink } from '@apollo/client/link/ws';
 import {
     DefaultContext,
     FetchResult,
@@ -56,9 +55,8 @@ export class WebSocketService {
     private readonly authService = inject(AuthenticationService);
 
     private readonly appConfigService = inject(AppConfigService);
-    private readonly subscriptionProtocol: 'graphql-ws' | 'transport-ws' = 'graphql-ws';
+    private wsLink!: GraphQLWsLink;
-    private wsLink: GraphQLWsLink | WebSocketLink;
+    private httpLinkHandler: HttpLinkHandler | undefined;
-    private httpLinkHandler: HttpLinkHandler;
 
     public getSubscription<T>(options: serviceOptions): Observable<FetchResult<T>> {
         const { apolloClientName, subscriptionOptions } = options;
@@ -110,8 +108,7 @@ export class WebSocketService {
             operation.setContext(({ headers }: DefaultContext) => ({
                 headers: {
                     ...headers,
-                    ...(this.subscriptionProtocol === 'graphql-ws' && { Authorization: `Bearer ${this.authService.getToken()}` }),
+                    Authorization: `Bearer ${this.authService.getToken()}`
-                    ...(this.subscriptionProtocol === 'transport-ws' && { 'X-Authorization': `Bearer ${this.authService.getToken()}` })
                 }
             }));
             return forward(operation);
@@ -120,8 +117,8 @@ export class WebSocketService {
         const errorLink = onError(({ graphQLErrors, networkError, operation, forward }) => {
             if (graphQLErrors) {
                 for (const error of graphQLErrors) {
-                    if (error.extensions && error.extensions['code'] === 'UNAUTHENTICATED') {
+                    if (error.extensions?.['code'] === 'UNAUTHENTICATED') {
-                        authLink(operation, forward);
+                        return authLink(operation, forward);
                     }
                 }
             }
@@ -129,6 +126,8 @@ export class WebSocketService {
             if (networkError) {
                 console.error(`[Network error]: ${networkError}`);
             }
+
+            return undefined;
         });
 
         const retryLink = new RetryLink({
@@ -145,8 +144,7 @@ export class WebSocketService {
 
         this.apollo.createNamed(options.apolloClientName, {
             headers: {
-                ...(this.subscriptionProtocol === 'graphql-ws' && { Authorization: `Bearer ${this.authService.getToken()}` }),
+                Authorization: `Bearer ${this.authService.getToken()}`
-                ...(this.subscriptionProtocol === 'transport-ws' && { 'X-Authorization': `Bearer ${this.authService.getToken()}` })
             },
             link: from([authLink, retryLink, errorLink, link]),
             cache: new InMemoryCache({ merge: true } as InMemoryCacheConfig)
@@ -155,22 +153,28 @@ export class WebSocketService {
 
     private createGraphQLWsLink(options: serviceOptions): void {
         this.wsLink = new GraphQLWsLink(
-            createClient({
+            this.createWsClient({
                 url: this.createWsUrl(options.wsUrl) + '/v2/ws/graphql',
                 connectionParams: () => ({
                     Authorization: 'Bearer ' + this.authService.getToken()
                 }),
                 on: {
-                    error: () => {
+                    error: () => this.reconnect(options)
-                        this.apollo.removeClient(options.apolloClientName);
-                        this.initSubscriptions(options);
-                    }
                 },
                 lazy: true
             })
         );
     }
 
+    protected createWsClient(clientOptions: ClientOptions): Client {
+        return createClient(clientOptions);
+    }
+
+    private reconnect(options: serviceOptions): void {
+        this.apollo.removeClient(options.apolloClientName);
+        this.initSubscriptions(options);
+    }
+
     private createHttpLinkHandler(options: serviceOptions): void {
         this.httpLinkHandler = options.httpUrl
             ? this.httpLink.create({