@webex/internal-plugin-mercury 3.11.0 → 3.12.0-llmrefactor.2

package/test/unit/spec/mercury.js

@@ -102,7 +102,48 @@ describe('plugin-mercury', () => {
       mercury = webex.internal.mercury;
     });
 
-    afterEach(() => {
+    afterEach(async () => {
+      // Clean up Mercury connections and internal state synchronously
+      // to avoid issues with fake timers
+      if (mercury) {
+        // Abort any pending backoff calls first
+        if (mercury.backoffCall) {
+          mercury.backoffCall.abort();
+          mercury.backoffCall = undefined;
+        }
+        if (mercury._shutdownSwitchoverBackoffCall) {
+          mercury._shutdownSwitchoverBackoffCall.abort();
+          mercury._shutdownSwitchoverBackoffCall = undefined;
+        }
+
+        // Close socket directly without waiting for offline event
+        if (mercury.socket) {
+          if (typeof mercury.socket.removeAllListeners === 'function') {
+            mercury.socket.removeAllListeners();
+          }
+          try {
+            mercury.socket.close();
+          } catch (e) {
+            // Ignore cleanup errors
+          }
+          mercury.socket = undefined;
+        }
+
+        // Clear any remaining connection promise
+        mercury._connectPromise = null;
+        mercury.connected = false;
+        mercury.connecting = false;
+      }
+
+      // Ensure mock socket is properly closed
+      if (mockWebSocket && typeof mockWebSocket.close === 'function') {
+        try {
+          mockWebSocket.close();
+        } catch (e) {
+          // Ignore cleanup errors
+        }
+      }
+
       if (socketOpenStub) {
         socketOpenStub.restore();
       }
@@ -110,6 +151,9 @@ describe('plugin-mercury', () => {
       if (Socket.getWebSocketConstructor.restore) {
         Socket.getWebSocketConstructor.restore();
       }
+
+      // Small delay to ensure all async operations complete
+      await new Promise((resolve) => setTimeout(resolve, 10));
     });
 
     describe('#listen()', () => {
@@ -268,18 +312,20 @@ describe('plugin-mercury', () => {
             .then(() => {
               assert.calledOnce(Socket.prototype.open);
 
-              return promiseTick(5);
+              // Wait for error handling from 1st failure to schedule 2nd timer
+              return promiseTick(10);
             })
             .then(() => {
               clock.tick(mercury.config.backoffTimeReset);
 
-              return promiseTick(5);
+              // Wait for 2nd socket.open call AND error handling to schedule 3rd timer
+              return promiseTick(10);
             })
             .then(() => {
               assert.calledTwice(Socket.prototype.open);
               clock.tick(2 * mercury.config.backoffTimeReset);
 
-              return promiseTick(5);
+              return promiseTick(10);
             })
             .then(() => {
               assert.calledThrice(Socket.prototype.open);
@@ -354,20 +400,20 @@ describe('plugin-mercury', () => {
             .then(() => {
               assert.calledOnce(Socket.prototype.open);
 
-              // I'm not sure why, but it's important the clock doesn't advance
-              // until a tick happens
-              return promiseTick(5);
+              // Wait for error handling from 1st failure to schedule 2nd timer
+              return promiseTick(10);
             })
             .then(() => {
               clock.tick(mercury.config.backoffTimeReset);
 
-              return promiseTick(5);
+              // Wait for 2nd socket.open call AND error handling to schedule 3rd timer
+              return promiseTick(10);
             })
             .then(() => {
               assert.calledTwice(Socket.prototype.open);
               clock.tick(2 * mercury.config.backoffTimeReset);
 
-              return promiseTick(5);
+              return promiseTick(10);
             })
             .then(() => {
               assert.calledThrice(Socket.prototype.open);
@@ -467,15 +513,49 @@ describe('plugin-mercury', () => {
         // });
 
         describe('when web-high-availability feature is enabled', () => {
-          it('marks current socket url as failed and get new one on Connection Error', () => {
+          it('marks the URL as failed when socket.open() fails with a ConnectionError', () => {
+            // When socket.open() fails, _prepareAndOpenSocket attaches webSocketUrl to the error.
+            // The error handler then uses wsUrl = newWSUrl || reason.webSocketUrl, so markFailedUrl
+            // is called with the URL obtained before the socket open attempt.
             webex.internal.feature.getFeature.returns(Promise.resolve(true));
             socketOpenStub.restore();
             socketOpenStub = sinon.stub(Socket.prototype, 'open').returns(Promise.resolve());
             socketOpenStub.onCall(0).returns(Promise.reject(new ConnectionError({code: 4001})));
             const promise = mercury.connect();
 
-            return promiseTick(7).then(() => {
+            return promiseTick(15).then(() => {
               assert.calledOnce(webex.internal.services.markFailedUrl);
+              assert.calledWith(
+                webex.internal.services.markFailedUrl,
+                sinon.match(/example-2\.com/)
+              );
+              clock.tick(1000);
+
+              return promise;
+            });
+          });
+
+          it('does not mark a URL as failed when no wsUrl is available in ConnectionError', () => {
+            // If the error has no webSocketUrl and newWSUrl is also undefined,
+            // markFailedUrl should not be called.
+            webex.internal.feature.getFeature.returns(Promise.resolve(true));
+
+            const errWithNoUrl = new ConnectionError({code: 4001});
+            // Explicitly ensure no webSocketUrl on the error
+            delete errWithNoUrl.webSocketUrl;
+
+            // Stub _prepareAndOpenSocket so it rejects without attaching webSocketUrl
+            // (simulating a failure before socket.open is reached, e.g. _prepareUrl failure)
+            const prepareAndOpenStub = sinon.stub(mercury, '_prepareAndOpenSocket');
+
+            prepareAndOpenStub.onCall(0).returns(Promise.reject(errWithNoUrl));
+            prepareAndOpenStub.returns(Promise.resolve('ws://example-2.com'));
+
+            const promise = mercury.connect();
+
+            return promiseTick(15).then(() => {
+              assert.notCalled(webex.internal.services.markFailedUrl);
+              prepareAndOpenStub.restore();
               clock.tick(1000);
 
               return promise;
@@ -499,17 +579,22 @@ describe('plugin-mercury', () => {
 
         // skipping due to apparent bug with lolex in all browsers but Chrome.
         skipInBrowser(it)('does not continue attempting to connect', () => {
-          mercury.connect();
+          const promise = mercury.connect();
 
-          return promiseTick(2)
-            .then(() => {
-              clock.tick(6 * webex.internal.mercury.config.backoffTimeReset);
+          // Wait for the connection to be established before proceeding
+          mockWebSocket.open();
 
-              return promiseTick(2);
-            })
-            .then(() => {
-              assert.calledOnce(Socket.prototype.open);
-            });
+          return promise.then(() =>
+            promiseTick(2)
+              .then(() => {
+                clock.tick(6 * webex.internal.mercury.config.backoffTimeReset);
+
+                return promiseTick(2);
+              })
+              .then(() => {
+                assert.calledOnce(Socket.prototype.open);
+              })
+          );
         });
       });
 
@@ -583,44 +668,6 @@ describe('plugin-mercury', () => {
       });
     });
 
-    describe('#logout()', () => {
-      it('calls disconnect and logs', () => {
-        sinon.stub(mercury.logger, 'info');
-        sinon.stub(mercury, 'disconnect');
-        mercury.logout();
-        assert.called(mercury.disconnect);
-        assert.calledTwice(mercury.logger.info);
-
-        assert.calledWith(mercury.logger.info.getCall(0), 'Mercury: logout() called');
-        assert.isTrue(
-          mercury.logger.info
-            .getCall(1)
-            .args[0].startsWith('Mercury: debug_mercury_logging stack: ')
-        );
-      });
-
-      it('uses the config.beforeLogoutOptionsCloseReason to disconnect and will send code 3050 for logout', () => {
-        sinon.stub(mercury, 'disconnect');
-        mercury.config.beforeLogoutOptionsCloseReason = 'done (permanent)';
-        mercury.logout();
-        assert.calledWith(mercury.disconnect, {code: 3050, reason: 'done (permanent)'});
-      });
-
-      it('uses the config.beforeLogoutOptionsCloseReason to disconnect and will send code 3050 for logout if the reason is different than standard', () => {
-        sinon.stub(mercury, 'disconnect');
-        mercury.config.beforeLogoutOptionsCloseReason = 'test';
-        mercury.logout();
-        assert.calledWith(mercury.disconnect, {code: 3050, reason: 'test'});
-      });
-
-      it('uses the config.beforeLogoutOptionsCloseReason to disconnect and will send undefined for logout if the reason is same as standard', () => {
-        sinon.stub(mercury, 'disconnect');
-        mercury.config.beforeLogoutOptionsCloseReason = 'done (forced)';
-        mercury.logout();
-        assert.calledWith(mercury.disconnect, undefined);
-      });
-    });
-
     describe('#disconnect()', () => {
       it('disconnects the WebSocket', () =>
         mercury
@@ -723,12 +770,15 @@ describe('plugin-mercury', () => {
           return promiseTick(webex.internal.mercury.config.backoffTimeReset).then(() => {
             // By this time backoffCall and mercury socket should be defined by the
             // 'connect' call
-            assert.isDefined(mercury.backoffCall, 'Mercury backoffCall is not defined');
+            assert.isDefined(
+              mercury.backoffCall,
+              'Mercury backoffCall is not defined'
+            );
             assert.isDefined(mercury.socket, 'Mercury socket is not defined');
             // Calling disconnect will abort the backoffCall, close the socket, and
             // reject the connect
             mercury.disconnect();
-            assert.isUndefined(mercury.backoffCall, 'Mercury backoffCall is still defined');
+            assert.notExists(mercury.backoffCall, 'Mercury backoffCall is still defined');
             // The socket will never be unset (which seems bad)
             assert.isDefined(mercury.socket, 'Mercury socket is not defined');
 
@@ -747,7 +797,8 @@ describe('plugin-mercury', () => {
           let reason;
 
           mercury.backoffCall = undefined;
-          mercury._attemptConnection('ws://example.com', (_reason) => {
+
+          const promise = mercury._attemptConnection('ws://example.com', (_reason) => {
             reason = _reason;
           });
 
@@ -756,6 +807,9 @@ describe('plugin-mercury', () => {
               reason.message,
               'Mercury: prevent socket open when backoffCall no longer defined'
             );
+
+            // Ensure the promise was actually rejected (short-circuited)
+            return assert.isRejected(promise);
           });
         });
 
@@ -1016,7 +1070,8 @@ describe('plugin-mercury', () => {
         });
 
         it('should be idempotent - no-op if already in progress', () => {
-          mercury._shutdownSwitchoverInProgress = true;
+          // Simulate an existing switchover in progress
+          mercury._shutdownSwitchoverBackoffCall = {abort: sinon.stub()};
 
           mercury._handleImminentShutdown();
 
@@ -1026,8 +1081,15 @@ describe('plugin-mercury', () => {
         it('should set switchover flags when called', () => {
           mercury._handleImminentShutdown();
 
-          assert.isTrue(mercury._shutdownSwitchoverInProgress);
           assert.isDefined(mercury._shutdownSwitchoverId);
+
+          assert.calledOnce(connectWithBackoffStub);
+          const callArgs = connectWithBackoffStub.firstCall.args;
+          assert.isUndefined(callArgs[0]); // webSocketUrl
+          assert.isObject(callArgs[1]); // context
+          assert.isTrue(callArgs[1].isShutdownSwitchover);
+          assert.isObject(callArgs[1].attemptOptions);
+          assert.isTrue(callArgs[1].attemptOptions.isShutdownSwitchover);
         });
 
         it('should call _connectWithBackoff with correct parameters', (done) => {
@@ -1039,6 +1101,8 @@ describe('plugin-mercury', () => {
             assert.isUndefined(callArgs[0]); // webSocketUrl
             assert.isObject(callArgs[1]); // context
             assert.isTrue(callArgs[1].isShutdownSwitchover);
+            assert.isObject(callArgs[1].attemptOptions);
+            assert.isTrue(callArgs[1].attemptOptions.isShutdownSwitchover);
             done();
           });
         });
@@ -1049,7 +1113,9 @@ describe('plugin-mercury', () => {
 
           mercury._handleImminentShutdown();
 
-          assert.isFalse(mercury._shutdownSwitchoverInProgress);
+          // When an exception happens synchronously, the backoff call should be cleared
+          assert.notExists(mercury._shutdownSwitchoverBackoffCall);
+          mercury._connectWithBackoff.restore();
         });
       });
 
@@ -1108,6 +1174,117 @@ describe('plugin-mercury', () => {
         });
       });
 
+      describe('#_onmessage() with missing data or eventType', () => {
+        beforeEach(() => {
+          sinon.stub(mercury, '_emit');
+          sinon.stub(mercury, '_setTimeOffset');
+          sinon.stub(mercury, '_applyOverrides');
+        });
+
+        afterEach(() => {
+          mercury._emit.restore();
+          mercury._setTimeOffset.restore();
+          mercury._applyOverrides.restore();
+        });
+
+        it('should emit generic event and return early when envelope.data (nested) is undefined', async () => {
+          const event = {
+            data: {
+              type: 'someType',
+              // no nested data property
+            },
+          };
+
+          const result = mercury._onmessage(event);
+
+          assert.instanceOf(result, Promise);
+          await result;
+          assert.calledOnceWithExactly(mercury._emit, 'event', event.data);
+          assert.notCalled(mercury._applyOverrides);
+        });
+
+        it('should emit generic event and return early when data.eventType is undefined', async () => {
+          const event = {
+            data: {
+              type: 'someType',
+              data: {
+                // no eventType property
+                someField: 'value',
+              },
+            },
+          };
+
+          const result = mercury._onmessage(event);
+
+          assert.instanceOf(result, Promise);
+          await result;
+          assert.calledOnceWithExactly(mercury._emit, 'event', event.data);
+          assert.notCalled(mercury._applyOverrides);
+        });
+
+        it('should emit generic event for messages without eventType (e.g. subscription responses)', async () => {
+          const event = {
+            data: {
+              id: 'msg-123',
+              sequenceNumber: 5,
+              data: {
+                statusCode: 200,
+              },
+            },
+          };
+
+          const result = mercury._onmessage(event);
+
+          assert.instanceOf(result, Promise);
+          await result;
+          assert.calledOnceWithExactly(mercury._emit, 'event', event.data);
+          assert.notCalled(mercury._applyOverrides);
+        });
+
+        it('should still process messages with a valid eventType', async () => {
+          const event = {
+            data: {
+              data: {
+                eventType: 'conversation.activity',
+              },
+            },
+          };
+
+          await mercury._onmessage(event);
+
+          // Normal flow emits namespace-specific events after processing handlers.
+          // The early-return guard only emits 'event', so asserting these proves the normal path was taken.
+          assert.calledWith(mercury._emit, 'event:conversation', event.data);
+          assert.calledWith(mercury._emit, 'event:conversation.activity', event.data);
+        });
+      });
+
+      describe('#_getEventHandlers()', () => {
+        it('should return an empty array when eventType is undefined', () => {
+          const result = mercury._getEventHandlers(undefined);
+
+          assert.deepEqual(result, []);
+        });
+
+        it('should return an empty array when eventType is null', () => {
+          const result = mercury._getEventHandlers(null);
+
+          assert.deepEqual(result, []);
+        });
+
+        it('should return an empty array when eventType is an empty string', () => {
+          const result = mercury._getEventHandlers('');
+
+          assert.deepEqual(result, []);
+        });
+
+        it('should return an empty array when namespace is not registered', () => {
+          const result = mercury._getEventHandlers('unknownNamespace.someEvent');
+
+          assert.deepEqual(result, []);
+        });
+      });
+
       describe('#_onclose() with code 4001 (shutdown replacement)', () => {
         let mockSocket, anotherSocket;
 
@@ -1201,8 +1378,6 @@ describe('plugin-mercury', () => {
           mercury._onclose(closeEvent, anotherSocket);
 
           // Verify listeners were removed from the old socket
-          // The _onclose method checks if sourceSocket !== this.socket (non-active)
-          // and then calls removeAllListeners in the else branch
           assert.calledOnce(anotherSocket.removeAllListeners);
         });
 
@@ -1220,6 +1395,69 @@ describe('plugin-mercury', () => {
         });
       });
 
+      describe('#_onclose() reconnect URL derivation', () => {
+        // The URL resolved and stored in _prepareAndOpenSocket (catalog-valid host).
+        const resolvedUrl = 'wss://mercury-connection.wbx2.com/v1/apps/wx2/registrations';
+        // The URL the native socket was actually opened with, after a lower
+        // layer (e.g. the same-site websocket proxy) rewrote it into a
+        // non-catalog host. Reusing this would fail host-catalog validation.
+        const proxiedUrl = 'wss://web.webex.com/webproxy/mercury-connection.wbx2.com/v1/apps/wx2';
+        let sourceSocket;
+
+        beforeEach(() => {
+          sourceSocket = {
+            url: proxiedUrl,
+            removeAllListeners: sinon.stub(),
+          };
+          mercury.socket = sourceSocket;
+          mercury.resolvedWebSocketUrl = resolvedUrl;
+          mercury.connected = true;
+          sinon.stub(mercury, '_emit');
+          sinon.stub(mercury, '_reconnect');
+          sinon.stub(mercury, 'unset');
+        });
+
+        afterEach(() => {
+          mercury._emit.restore();
+          mercury._reconnect.restore();
+          mercury.unset.restore();
+        });
+
+        [1001, 1005, 1006, 1011].forEach((code) => {
+          it(`reconnects an active socket using the stored resolved URL (not the proxied socket.url) for code ${code}`, () => {
+            mercury._onclose({code, reason: 'transient'}, sourceSocket);
+
+            assert.calledOnceWithExactly(mercury._reconnect, resolvedUrl);
+            assert.neverCalledWith(mercury._reconnect, proxiedUrl);
+          });
+        });
+
+        it('reconnects using the stored resolved URL for a normal close (1000) with a reconnect reason', () => {
+          mercury._onclose({code: 1000, reason: 'done (forced)'}, sourceSocket);
+
+          assert.calledOnceWithExactly(mercury._reconnect, resolvedUrl);
+        });
+
+        it('does not reconnect when a non-active (old) socket closes', () => {
+          const oldSocket = {
+            url: proxiedUrl,
+            removeAllListeners: sinon.stub(),
+          };
+
+          mercury._onclose({code: 1006, reason: 'transient'}, oldSocket);
+
+          assert.notCalled(mercury._reconnect);
+        });
+
+        it('falls back to socket.url when resolvedWebSocketUrl is not set', () => {
+          mercury.resolvedWebSocketUrl = undefined;
+
+          mercury._onclose({code: 1006, reason: 'transient'}, sourceSocket);
+
+          assert.calledOnceWithExactly(mercury._reconnect, proxiedUrl);
+        });
+      });
+
       describe('shutdown switchover with retry logic', () => {
         let connectWithBackoffStub;
 
@@ -1243,13 +1481,12 @@ describe('plugin-mercury', () => {
 
           mercury._handleImminentShutdown();
 
-          // Give it a tick for the async call to happen
           process.nextTick(() => {
             assert.calledOnce(connectWithBackoffStub);
             const callArgs = connectWithBackoffStub.firstCall.args;
 
-            assert.isUndefined(callArgs[0]); // webSocketUrl is undefined
-            assert.isObject(callArgs[1]); // context object
+            assert.isUndefined(callArgs[0]); // webSocketUrl
+            assert.isObject(callArgs[1]);
             assert.isTrue(callArgs[1].isShutdownSwitchover);
             assert.isObject(callArgs[1].attemptOptions);
             assert.isTrue(callArgs[1].attemptOptions.isShutdownSwitchover);
@@ -1257,18 +1494,21 @@ describe('plugin-mercury', () => {
           });
         });
 
-        it('should set _shutdownSwitchoverInProgress flag during switchover', () => {
-          connectWithBackoffStub.returns(new Promise(() => {})); // Never resolves
+        it('should set _shutdownSwitchoverBackoffCall during switchover', () => {
+          // Since _connectWithBackoff is stubbed, we can just check that
+          // _handleImminentShutdown sets the flag
+          connectWithBackoffStub.callsFake(() => {
+            return new Promise(() => {}); // Never resolves
+          });
 
           mercury._handleImminentShutdown();
 
-          assert.isTrue(mercury._shutdownSwitchoverInProgress);
+          // The switchover id is set when switchover starts
+          assert.isDefined(mercury._shutdownSwitchoverId);
         });
 
         it('should emit success event when switchover completes', async () => {
-          // We need to actually call the onSuccess callback to trigger the event
           connectWithBackoffStub.callsFake((url, context) => {
-            // Simulate successful connection by calling onSuccess
             if (context && context.attemptOptions && context.attemptOptions.onSuccess) {
               const mockSocket = {url: 'ws://new-socket.com'};
               context.attemptOptions.onSuccess(mockSocket, 'ws://new-socket.com');
@@ -1278,7 +1518,6 @@ describe('plugin-mercury', () => {
 
           mercury._handleImminentShutdown();
 
-          // Wait for async operations
           await promiseTick(50);
 
           const emitCalls = mercury._emit.getCalls();
@@ -1292,12 +1531,11 @@ describe('plugin-mercury', () => {
         it('should emit failure event when switchover exhausts retries', async () => {
           const testError = new Error('Connection failed');
 
-          connectWithBackoffStub.returns(Promise.reject(testError));
+          connectWithBackoffStub.callsFake(() => Promise.reject(testError));
 
           mercury._handleImminentShutdown();
           await promiseTick(50);
 
-          // Check if failure event was emitted
           const emitCalls = mercury._emit.getCalls();
           const hasFailureEvent = emitCalls.some(
             (call) =>
@@ -1310,12 +1548,11 @@ describe('plugin-mercury', () => {
         });
 
         it('should allow old socket to be closed by server after switchover failure', async () => {
-          connectWithBackoffStub.returns(Promise.reject(new Error('Failed')));
+          connectWithBackoffStub.callsFake(() => Promise.reject(new Error('Failed')));
 
           mercury._handleImminentShutdown();
           await promiseTick(50);
 
-          // Old socket should not be closed immediately - server will close it
           assert.equal(mercury.socket.removeAllListeners.callCount, 0);
         });
       });
@@ -1399,6 +1636,23 @@ describe('plugin-mercury', () => {
           assert.equal(result, 'ws://example.com');
         });
 
+        it('stores the resolved (pre-proxy) URL in resolvedWebSocketUrl', async () => {
+          await mercury._prepareAndOpenSocket(mockSocket, 'ws://original.com', false);
+
+          // The stored value is the URL resolved by _prepareUrl, not the raw
+          // input, so _onclose can later reconnect from a catalog-valid host.
+          assert.equal(mercury.resolvedWebSocketUrl, 'ws://example.com');
+        });
+
+        it('does not store resolvedWebSocketUrl for shutdown switchover connections', async () => {
+          mercury.resolvedWebSocketUrl = 'ws://existing.com';
+
+          await mercury._prepareAndOpenSocket(mockSocket, 'ws://original.com', true);
+
+          // Shutdown switchover should not overwrite the stored URL
+          assert.equal(mercury.resolvedWebSocketUrl, 'ws://existing.com');
+        });
+
         it('should handle errors during socket open', async () => {
           mockSocket.open.returns(Promise.reject(new Error('Open failed')));
 
@@ -1407,23 +1661,21 @@ describe('plugin-mercury', () => {
             assert.fail('Should have thrown an error');
           } catch (err) {
             assert.equal(err.message, 'Open failed');
+            assert.equal(err.webSocketUrl, 'ws://example.com');
           }
         });
       });
 
       describe('#_attemptConnection() with shutdown switchover', () => {
-        let mockSocket, prepareAndOpenSocketStub, callback;
+        let prepareAndOpenSocketStub, callback;
 
         beforeEach(() => {
-          mockSocket = {
-            url: 'ws://test.com',
-          };
           prepareAndOpenSocketStub = sinon
             .stub(mercury, '_prepareAndOpenSocket')
             .returns(Promise.resolve('ws://new-socket.com'));
           callback = sinon.stub();
-          mercury._shutdownSwitchoverBackoffCall = {}; // Mock backoff call
-          mercury.socket = mockSocket;
+          mercury._shutdownSwitchoverBackoffCall = {abort: sinon.stub()};
+          mercury.socket = {url: 'ws://test.com'};
           mercury.connected = true;
           sinon.stub(mercury, '_emit');
           sinon.stub(mercury, '_attachSocketEventListeners');
@@ -1433,6 +1685,7 @@ describe('plugin-mercury', () => {
           prepareAndOpenSocketStub.restore();
           mercury._emit.restore();
           mercury._attachSocketEventListeners.restore();
+          mercury._shutdownSwitchoverBackoffCall = undefined;
         });
 
         it('should not set socket reference before opening for shutdown switchover', async () => {
@@ -1441,13 +1694,10 @@ describe('plugin-mercury', () => {
           await mercury._attemptConnection('ws://test.com', callback, {
             isShutdownSwitchover: true,
             onSuccess: (newSocket, url) => {
-              // During onSuccess, verify original socket is still set
-              // (socket swap happens inside onSuccess callback in _handleImminentShutdown)
               assert.equal(mercury.socket, originalSocket);
             },
           });
 
-          // After onSuccess, socket should still be original since we only swap in _handleImminentShutdown
           assert.equal(mercury.socket, originalSocket);
         });
 
@@ -1464,12 +1714,9 @@ describe('plugin-mercury', () => {
         });
 
         it('should emit shutdown switchover complete event', async () => {
-          const oldSocket = mercury.socket;
-
           await mercury._attemptConnection('ws://test.com', callback, {
             isShutdownSwitchover: true,
             onSuccess: (newSocket, url) => {
-              // Simulate the onSuccess callback behavior
               mercury.socket = newSocket;
               mercury.connected = true;
               mercury._emit('event:mercury_shutdown_switchover_complete', {url});
@@ -1486,15 +1733,12 @@ describe('plugin-mercury', () => {
         it('should use simpler error handling for shutdown switchover failures', async () => {
           prepareAndOpenSocketStub.returns(Promise.reject(new Error('Connection failed')));
 
-          try {
-            await mercury._attemptConnection('ws://test.com', callback, {
+          await mercury
+            ._attemptConnection('ws://test.com', callback, {
               isShutdownSwitchover: true,
-            });
-          } catch (err) {
-            // Error should be caught and passed to callback
-          }
+            })
+            .catch(() => {});
 
-          // Should call callback with error for retry
           assert.calledOnce(callback);
           assert.instanceOf(callback.firstCall.args[0], Error);
         });
@@ -1514,35 +1758,26 @@ describe('plugin-mercury', () => {
       });
 
       describe('#_connectWithBackoff() with shutdown switchover', () => {
-        // Note: These tests verify the parameterization logic without running real backoff timers
-        // to avoid test hangs. The backoff mechanism itself is tested in other test suites.
-
         it('should use shutdown-specific parameters when called', () => {
-          // Stub _connectWithBackoff to prevent real execution
           const connectWithBackoffStub = sinon
             .stub(mercury, '_connectWithBackoff')
             .returns(Promise.resolve());
 
           mercury._handleImminentShutdown();
 
-          // Verify it was called with shutdown context
           assert.calledOnce(connectWithBackoffStub);
           const callArgs = connectWithBackoffStub.firstCall.args;
-          assert.isObject(callArgs[1]); // context
+          // First arg is the URL, second is the context object
+          assert.isObject(callArgs[1]);
           assert.isTrue(callArgs[1].isShutdownSwitchover);
 
           connectWithBackoffStub.restore();
         });
 
         it('should pass shutdown switchover options to _attemptConnection', () => {
-          // Stub _attemptConnection to verify it receives correct options
           const attemptStub = sinon.stub(mercury, '_attemptConnection');
-          attemptStub.callsFake((url, callback) => {
-            // Immediately succeed
-            callback();
-          });
+          attemptStub.callsFake((url, cb) => cb());
 
-          // Call _connectWithBackoff with shutdown context
           const context = {
             isShutdownSwitchover: true,
             attemptOptions: {
@@ -1551,25 +1786,25 @@ describe('plugin-mercury', () => {
             },
           };
 
-          // Start the backoff
           const promise = mercury._connectWithBackoff(undefined, context);
 
-          // Check that _attemptConnection was called with shutdown options
           return promise.then(() => {
             assert.calledOnce(attemptStub);
             const callArgs = attemptStub.firstCall.args;
-            assert.isObject(callArgs[2]); // options parameter
+            assert.isObject(callArgs[2]);
             assert.isTrue(callArgs[2].isShutdownSwitchover);
-
             attemptStub.restore();
           });
         });
 
-        it('should set and clear state flags appropriately', () => {
-          // Stub to prevent actual connection
-          sinon.stub(mercury, '_attemptConnection').callsFake((url, callback) => callback());
+        it('should set _shutdownSwitchoverBackoffCall during attempt and clear it after completion', () => {
+          let backoffCallDuringAttempt;
 
-          mercury._shutdownSwitchoverInProgress = true;
+          sinon.stub(mercury, '_attemptConnection').callsFake((url, cb) => {
+            // Capture the backoff call state during the connection attempt
+            backoffCallDuringAttempt = mercury._shutdownSwitchoverBackoffCall;
+            cb();
+          });
 
           const promise = mercury._connectWithBackoff(undefined, {
             isShutdownSwitchover: true,
@@ -1577,8 +1812,13 @@ describe('plugin-mercury', () => {
           });
 
           return promise.then(() => {
-            // Should be cleared after completion
-            assert.isFalse(mercury._shutdownSwitchoverInProgress);
+            // Verify backoff call was set during the operation
+            assert.isDefined(backoffCallDuringAttempt, 'backoffCall should be set during attempt');
+            // Verify backoff call is cleared after completion
+            assert.notExists(
+              mercury._shutdownSwitchoverBackoffCall,
+              'backoffCall should be cleared after completion'
+            );
             mercury._attemptConnection.restore();
           });
         });
@@ -1588,30 +1828,28 @@ describe('plugin-mercury', () => {
         let abortStub;
 
         beforeEach(() => {
-          mercury.socket = {
-            close: sinon.stub().returns(Promise.resolve()),
-            removeAllListeners: sinon.stub(),
-          };
           abortStub = sinon.stub();
-          mercury._shutdownSwitchoverBackoffCall = {
-            abort: abortStub,
-          };
+          mercury._shutdownSwitchoverBackoffCall = {abort: abortStub};
         });
 
-        it('should abort shutdown switchover backoff call on disconnect', async () => {
-          await mercury.disconnect();
+        it('should abort shutdown switchover backoff call on disconnect', () => {
+          // Don't set up socket so disconnect resolves immediately
+          mercury.socket = undefined;
 
-          assert.calledOnce(abortStub);
+          return mercury.disconnect().then(() => {
+            assert.calledOnce(abortStub);
+          });
         });
 
-        it('should handle disconnect when no switchover is in progress', async () => {
+        it('should handle disconnect when no switchover is in progress', () => {
           mercury._shutdownSwitchoverBackoffCall = undefined;
+          // Don't set up socket so disconnect resolves immediately
+          mercury.socket = undefined;
 
-          // Should not throw
-          await mercury.disconnect();
-
-          // Should still close the socket
-          assert.calledOnce(mercury.socket.close);
+          return mercury.disconnect().then(() => {
+            // If we get here, disconnect completed without error
+            assert.notExists(mercury._shutdownSwitchoverBackoffCall);
+          });
         });
       });
     });