From c1c4d595569116d11340f176601cfb1b6e3d59dd Mon Sep 17 00:00:00 2001
From: Bailey Pearson <bailey.pearson@mongodb.com>
Date: Fri, 26 Aug 2022 13:49:12 -0400
Subject: [PATCH 1/2] fix: only set loadBalanced on handshake when explicitly
 set

---
 src/cmap/connect.ts | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/src/cmap/connect.ts b/src/cmap/connect.ts
index 354e1043ec..a94a0c0f33 100644
--- a/src/cmap/connect.ts
+++ b/src/cmap/connect.ts
@@ -214,7 +214,7 @@ export interface HandshakeDocument extends Document {
   client: ClientMetadata;
   compression: string[];
   saslSupportedMechs?: string;
-  loadBalanced: boolean;
+  loadBalanced?: boolean;
 }
 
 function prepareHandshakeDocument(authContext: AuthContext, callback: Callback<HandshakeDocument>) {
@@ -226,10 +226,13 @@ function prepareHandshakeDocument(authContext: AuthContext, callback: Callback<H
     [serverApi?.version ? 'hello' : LEGACY_HELLO_COMMAND]: true,
     helloOk: true,
     client: options.metadata || makeClientMetadata(options),
-    compression: compressors,
-    loadBalanced: options.loadBalanced
+    compression: compressors
   };
 
+  if (typeof options.loadBalanced === 'boolean' && Boolean(options.loadBalanced) === true) {
+    handshakeDoc.loadBalanced = true;
+  }
+
   const credentials = authContext.credentials;
   if (credentials) {
     if (credentials.mechanism === AuthMechanism.MONGODB_DEFAULT && credentials.username) {

From be5d979d2cbc98c6c6b8afe335bc4fd7379f8e45 Mon Sep 17 00:00:00 2001
From: Bailey Pearson <bailey.pearson@mongodb.com>
Date: Fri, 26 Aug 2022 15:19:01 -0400
Subject: [PATCH 2/2] test: add unit tests for loadBalanced handshake fix

---
 src/cmap/connect.ts            | 12 ++++++--
 test/unit/cmap/connect.test.js | 52 +++++++++++++++++++++++++++++++++-
 2 files changed, 61 insertions(+), 3 deletions(-)

diff --git a/src/cmap/connect.ts b/src/cmap/connect.ts
index a94a0c0f33..fedf731132 100644
--- a/src/cmap/connect.ts
+++ b/src/cmap/connect.ts
@@ -217,7 +217,15 @@ export interface HandshakeDocument extends Document {
   loadBalanced?: boolean;
 }
 
-function prepareHandshakeDocument(authContext: AuthContext, callback: Callback<HandshakeDocument>) {
+/**
+ * @internal
+ *
+ * This function is only exposed for testing purposes.
+ */
+export function prepareHandshakeDocument(
+  authContext: AuthContext,
+  callback: Callback<HandshakeDocument>
+) {
   const options = authContext.options;
   const compressors = options.compressors ? options.compressors : [];
   const { serverApi } = authContext.connection;
@@ -229,7 +237,7 @@ function prepareHandshakeDocument(authContext: AuthContext, callback: Callback<H
     compression: compressors
   };
 
-  if (typeof options.loadBalanced === 'boolean' && Boolean(options.loadBalanced) === true) {
+  if (options.loadBalanced === true) {
     handshakeDoc.loadBalanced = true;
   }
 
diff --git a/test/unit/cmap/connect.test.js b/test/unit/cmap/connect.test.js
index c4fab156c6..ce8e0e65b5 100644
--- a/test/unit/cmap/connect.test.js
+++ b/test/unit/cmap/connect.test.js
@@ -5,12 +5,16 @@ const { expect } = require('chai');
 const EventEmitter = require('events');
 const { setTimeout } = require('timers');
 
-const { connect } = require('../../../src/cmap/connect');
+const {
+  connect,
+  prepareHandshakeDocument: prepareHandshakeDocumentCb
+} = require('../../../src/cmap/connect');
 const { MongoCredentials } = require('../../../src/cmap/auth/mongo_credentials');
 const { genClusterTime } = require('../../tools/common');
 const { MongoNetworkError } = require('../../../src/error');
 const { HostAddress, isHello } = require('../../../src/utils');
 const { LEGACY_HELLO_COMMAND } = require('../../../src/constants');
+const { promisify } = require('util');
 
 describe('Connect Tests', function () {
   const test = {};
@@ -112,4 +116,50 @@ describe('Connect Tests', function () {
       done();
     });
   }).skipReason = 'TODO(NODE-2941): stop using 240.0.0.1 in tests';
+
+  context('prepareHandshakeDocument', () => {
+    const prepareHandshakeDocument = promisify(prepareHandshakeDocumentCb);
+
+    context('loadBalanced option', () => {
+      context('when loadBalanced is not set as an option', () => {
+        it('does not set loadBalanced on the handshake document', async () => {
+          const options = {};
+          const authContext = {
+            connection: {},
+            options
+          };
+          const handshakeDocument = await prepareHandshakeDocument(authContext);
+          expect(handshakeDocument).not.to.have.property('loadBalanced');
+        });
+      });
+
+      context('when loadBalanced is set to false', () => {
+        it('does not set loadBalanced on the handshake document', async () => {
+          const options = {
+            loadBalanced: false
+          };
+          const authContext = {
+            connection: {},
+            options
+          };
+          const handshakeDocument = await prepareHandshakeDocument(authContext);
+          expect(handshakeDocument).not.to.have.property('loadBalanced');
+        });
+      });
+
+      context('when loadBalanced is set to true', () => {
+        it('does set loadBalanced on the handshake document', async () => {
+          const options = {
+            loadBalanced: true
+          };
+          const authContext = {
+            connection: {},
+            options
+          };
+          const handshakeDocument = await prepareHandshakeDocument(authContext);
+          expect(handshakeDocument).to.have.property('loadBalanced', true);
+        });
+      });
+    });
+  });
 });