fixup: remove atlas restrictive check, add unit tests

Author: Anemy

packages/compass-connections/src/stores/connections-store-redux.spec.tsx

@@ -10,6 +10,8 @@ import {
 } from '@mongodb-js/testing-library-compass';
 import React from 'react';
 import { InMemoryConnectionStorage } from '@mongodb-js/connection-storage/provider';
+import { getDataServiceForConnection } from './connections-store-redux';
+import { type ConnectionInfo } from '@mongodb-js/connection-info';
 
 const mockConnections = [
   {
@@ -34,6 +36,13 @@ const mockConnections = [
   },
 ];
 
+const connectionInfoWithAtlasMetadata = {
+  ...createDefaultConnectionInfo(),
+  atlasMetadata: {
+    clusterName: 'pineapple',
+  } as ConnectionInfo['atlasMetadata'],
+};
+
 function renderCompassConnections(opts?: RenderConnectionsOptions) {
   return render(
     <div>
@@ -274,6 +283,101 @@ describe('CompassConnections store', function () {
         await connectionStorage.load({ id: mockConnections[0].id })
       ).to.have.nested.property('favorite.name', 'turtles');
     });
+
+    it('should ignore server heartbeat failed events that are not non-retryable error codes', async function () {
+      const { connectionsStore } = renderCompassConnections({
+        connectFn: async () => {
+          await wait(1);
+          return {};
+        },
+      });
+
+      // Wait till we're connected.
+      await connectionsStore.actions.connect(connectionInfoWithAtlasMetadata);
+
+      const connections = connectionsStore.getState().connections;
+      expect(connections.ids).to.have.lengthOf(1);
+
+      const dataService = getDataServiceForConnection(
+        connectionInfoWithAtlasMetadata.id
+      );
+
+      let didDisconnect = false;
+      let didCheckForConnected = false;
+      sinon.stub(dataService, 'disconnect').callsFake(async () => {
+        didDisconnect = true;
+        return Promise.resolve();
+      });
+      dataService.isConnected = () => {
+        // If this is called we know the error wasn't handled properly.
+        didCheckForConnected = true;
+        return true;
+      };
+
+      let didReceiveCallToHeartbeatFailedListener = false;
+      dataService.on('serverHeartbeatFailed', () => {
+        didReceiveCallToHeartbeatFailedListener = true;
+      });
+
+      // Send a heartbeat fail with an error that's not a non-retryable error code.
+      dataService['emit']('serverHeartbeatFailed', {
+        failure: new Error('code: 1234, Not the error we are looking for'),
+      });
+
+      // Wait for the listener to handle the message.
+      await waitFor(() => {
+        expect(didReceiveCallToHeartbeatFailedListener).to.be.true;
+      });
+      await wait(1);
+
+      expect(didDisconnect).to.be.false;
+      expect(didCheckForConnected).to.be.false;
+    });
+
+    it('should listen for non-retryable errors on server heartbeat failed events and disconnect the data service when encountered', async function () {
+      const { connectionsStore } = renderCompassConnections({
+        connectFn: async () => {
+          await wait(1);
+          return {};
+        },
+      });
+
+      // Wait till we're connected.
+      await connectionsStore.actions.connect(connectionInfoWithAtlasMetadata);
+
+      const connections = connectionsStore.getState().connections;
+      expect(connections.ids).to.have.lengthOf(1);
+
+      const dataService = getDataServiceForConnection(
+        connectionInfoWithAtlasMetadata.id
+      );
+
+      let didDisconnect = false;
+      sinon.stub(dataService, 'disconnect').callsFake(async () => {
+        didDisconnect = true;
+        return Promise.resolve();
+      });
+      dataService.isConnected = () => true;
+
+      // Send a heartbeat fail with an error that's a non-retryable error code.
+      dataService['emit']('serverHeartbeatFailed', {
+        failure: new Error('code: 3003, reason: Insufficient permissions'),
+      });
+
+      await waitFor(() => {
+        expect(didDisconnect).to.be.true;
+      });
+
+      await waitFor(function () {
+        const titleNode = screen.getByText('Unable to connect to pineapple');
+        expect(titleNode).to.be.visible;
+
+        const descriptionNode = screen.getByText(
+          'Reason: Insufficient permissions. To use continue to use this connection either disconnect and reconnect, or refresh your page.'
+        );
+        expect(descriptionNode).to.be.visible;
+      });
+    });
   });
 
   describe('#saveAndConnect', function () {

packages/compass-connections/src/stores/connections-store-redux.ts

@@ -1466,6 +1466,15 @@ function isAtlasStreamsInstance(
   }
 }
 
+// We listen for non-retry-able errors on failed server heartbeats.
+// These can happen when:
+// - A user's session has ended.
+// - The user's roles have changed.
+// - The cluster / group they are trying to connect to has since been deleted.
+// When we encounter one we disconnect. This is to avoid polluting logs/metrics
+// and to avoid constantly retrying to connect when we know it'll fail.
+// These error codes can be found at
+// https://github.com/10gen/mms/blob/de2a9c463cfe530efb8e2a0941033e8207b6cb11/server/src/main/com/xgen/cloud/services/clusterconnection/runtime/res/CustomCloseCodes.java
 const NonRetryableErrorCodes = [3000, 3003, 4004, 1008] as const;
 const NonRetryableErrorDescriptionFallbacks: {
   [code in typeof NonRetryableErrorCodes[number]]: string;
@@ -1501,7 +1510,9 @@ const openConnectionClosedWithNonRetryableErrorToast = (
 ) => {
   openToast(`non-retryable-error-encountered--${connectionInfo.id}`, {
     title: `Unable to connect to ${getConnectionTitle(connectionInfo)}`,
-    description: `Reason: ${getDescriptionForNonRetryableError(error)}`,
+    description: `Reason: ${getDescriptionForNonRetryableError(
+      error
+    )}. To use continue to use this connection either disconnect and reconnect, or refresh your page.`,
     variant: 'warning',
   });
 };
@@ -1701,34 +1712,32 @@ const connectWithOptions = (
         }
 
         let showedNonRetryableErrorToast = false;
-        if (connectionInfo.atlasMetadata) {
-          // When we're on cloud we listen for non-retry-able errors on failed server
-          // heartbeats. These can happen when:
-          // - A user's session has ended.
-          // - The user's roles have changed.
-          // - The cluster / group they are trying to connect to has since been deleted.
-          // When we encounter one we disconnect. This is to avoid polluting logs/metrics
-          // and to avoid constantly retrying to connect when we know it'll fail.
-          dataService.on(
-            'serverHeartbeatFailed',
-            (evt: ServerHeartbeatFailedEvent) => {
-              if (!isNonRetryableHeartbeatFailure(evt)) {
-                return;
-              }
-
-              if (!dataService.isConnected() || showedNonRetryableErrorToast) {
-                return;
-              }
+        // Listen for non-retry-able errors on failed server heartbeats.
+        // These can happen when:
+        // - A user's session has ended.
+        // - The user's roles have changed.
+        // - The cluster / group they are trying to connect to has since been deleted.
+        // When we encounter one we disconnect. This is to avoid polluting logs/metrics
+        // and to avoid constantly retrying to connect when we know it'll fail.
+        dataService.on(
+          'serverHeartbeatFailed',
+          (evt: ServerHeartbeatFailedEvent) => {
+            if (!isNonRetryableHeartbeatFailure(evt)) {
+              return;
+            }
 
-              openConnectionClosedWithNonRetryableErrorToast(
-                connectionInfo,
-                evt.failure
-              );
-              showedNonRetryableErrorToast = true;
-              void dataService.disconnect();
+            if (!dataService.isConnected() || showedNonRetryableErrorToast) {
+              return;
             }
-          );
-        }
+
+            openConnectionClosedWithNonRetryableErrorToast(
+              connectionInfo,
+              evt.failure
+            );
+            showedNonRetryableErrorToast = true;
+            void dataService.disconnect();
+          }
+        );
 
         dataService.on('oidcAuthFailed', (error) => {
           openToast('oidc-auth-failed', {