RUST-1163 Fix load balancer tests (#576)

Author: patrickfreed

.evergreen/config.yml

@@ -349,6 +349,7 @@ functions:
             AUTH=${AUTH} \
             SSL=${SSL} \
             REQUIRE_API_VERSION=${REQUIRE_API_VERSION} \
+            LOAD_BALANCER=${LOAD_BALANCER} \
             sh ${DRIVERS_TOOLS}/.evergreen/run-orchestration.sh
     # run-orchestration generates expansion file with the MONGODB_URI for the cluster
     - command: expansions.update
@@ -883,6 +884,7 @@ tasks:
         vars:
           MONGODB_VERSION: "latest"
           TOPOLOGY: "sharded_cluster"
+          LOAD_BALANCER: "true"
       - func: "start load balancer"
       - func: "run tests"
 

src/client/options/mod.rs

@@ -583,9 +583,6 @@ pub(crate) struct TestOptions {
 
     /// Mock response for `SrvPollingMonitor::lookup_hosts`.
     pub(crate) mock_lookup_hosts: Option<Result<LookupHosts>>,
-
-    /// Mock the `serviceId` response for a load-balanced hello.
-    pub(crate) mock_service_id: bool,
 }
 
 fn default_hosts() -> Vec<ServerAddress> {

src/cmap/conn/wire/test.rs

@@ -6,13 +6,14 @@ use crate::{
     bson::{doc, Bson},
     cmap::options::StreamOptions,
     runtime::AsyncStream,
-    test::{CLIENT_OPTIONS, LOCK},
+    test::{log_uncaptured, CLIENT_OPTIONS, LOCK},
 };
 
 #[cfg_attr(feature = "tokio-runtime", tokio::test)]
 #[cfg_attr(feature = "async-std-runtime", async_std::test)]
 async fn basic() {
-    if CLIENT_OPTIONS.tls_options().is_some() {
+    if CLIENT_OPTIONS.tls_options().is_some() || CLIENT_OPTIONS.load_balanced.unwrap_or(false) {
+        log_uncaptured("skipping conn::wire:basic test due to TLS and/or load balanced topology");
         return;
     }
 

src/cmap/establish/handshake/mod.rs

@@ -144,8 +144,6 @@ pub(crate) struct Handshaker {
     /// given the same pool options, so it can be created at the time the Handshaker is created.
     command: Command,
     credential: Option<Credential>,
-    #[cfg(test)]
-    mock_service_id: bool,
     // This field is not read without a compression feature flag turned on.
     #[allow(dead_code)]
     compressors: Option<Vec<Compressor>>,
@@ -162,9 +160,6 @@ impl Handshaker {
         let mut command =
             is_master_command(options.as_ref().and_then(|opts| opts.server_api.as_ref()));
 
-        #[cfg(test)]
-        let mut mock_service_id = false;
-
         if let Some(options) = options {
             if let Some(app_name) = options.app_name {
                 metadata.application = Some(AppMetadata { name: app_name });
@@ -196,10 +191,7 @@ impl Handshaker {
             if options.load_balanced {
                 command.body.insert("loadBalanced", true);
             }
-            #[cfg(test)]
-            {
-                mock_service_id = options.mock_service_id;
-            }
+
             // Add compressors to handshake.
             // See https://github.com/mongodb/specifications/blob/master/source/compression/OP_COMPRESSED.rst
             if let Some(ref compressors) = options.compressors {
@@ -219,8 +211,6 @@ impl Handshaker {
         Self {
             command,
             credential,
-            #[cfg(test)]
-            mock_service_id,
             compressors,
         }
     }
@@ -237,24 +227,7 @@ impl Handshaker {
         let client_first = set_speculative_auth_info(&mut command.body, self.credential.as_ref())?;
 
         let mut is_master_reply = run_is_master(conn, command, topology, handler).await?;
-        // TODO PM-2369 Remove serviceId mocking when it's returned by the server.
-        #[cfg(test)]
-        {
-            if self.command.body.contains_key("loadBalanced")
-                && is_master_reply.command_response.service_id.is_none()
-                && self.mock_service_id
-            {
-                is_master_reply.command_response.service_id = Some(
-                    is_master_reply
-                        .command_response
-                        .topology_version
-                        .as_ref()
-                        .unwrap()
-                        .get_object_id("processId")
-                        .unwrap(),
-                );
-            }
-        }
+
         if self.command.body.contains_key("loadBalanced")
             && is_master_reply.command_response.service_id.is_none()
         {
@@ -326,8 +299,6 @@ pub(crate) struct HandshakerOptions {
     driver_info: Option<DriverInfo>,
     server_api: Option<ServerApi>,
     load_balanced: bool,
-    #[cfg(test)]
-    mock_service_id: bool,
 }
 
 impl From<ConnectionPoolOptions> for HandshakerOptions {
@@ -339,8 +310,6 @@ impl From<ConnectionPoolOptions> for HandshakerOptions {
             driver_info: options.driver_info,
             server_api: options.server_api,
             load_balanced: options.load_balanced.unwrap_or(false),
-            #[cfg(test)]
-            mock_service_id: options.mock_service_id,
         }
     }
 }
@@ -354,8 +323,6 @@ impl From<ClientOptions> for HandshakerOptions {
             driver_info: options.driver_info,
             server_api: options.server_api,
             load_balanced: options.load_balanced.unwrap_or(false),
-            #[cfg(test)]
-            mock_service_id: options.test_options.map_or(false, |to| to.mock_service_id),
         }
     }
 }

src/cmap/options.rs

@@ -96,11 +96,6 @@ pub(crate) struct ConnectionPoolOptions {
 
     /// Whether or not the client is connecting to a MongoDB cluster through a load balancer.
     pub(crate) load_balanced: Option<bool>,
-
-    /// Whether or not to mock the `serviceId` field of a hello through a load balancer.
-    #[cfg(test)]
-    #[serde(skip)]
-    pub(crate) mock_service_id: bool,
 }
 
 impl ConnectionPoolOptions {
@@ -122,11 +117,6 @@ impl ConnectionPoolOptions {
             #[cfg(test)]
             ready: None,
             load_balanced: options.load_balanced,
-            #[cfg(test)]
-            mock_service_id: options
-                .test_options
-                .as_ref()
-                .map_or(false, |to| to.mock_service_id),
         }
     }
 

src/test/mod.rs

@@ -90,12 +90,6 @@ pub(crate) fn update_options_for_testing(options: &mut ClientOptions) {
     if options.server_api.is_none() {
         options.server_api = SERVER_API.clone();
     }
-    if LOAD_BALANCED_SINGLE_URI
-        .as_ref()
-        .map_or(false, |uri| !uri.is_empty())
-    {
-        options.test_options_mut().mock_service_id = true;
-    }
     if options.compressors.is_none() {
         options.compressors = get_compressors();
     }

src/test/spec/connection_stepdown.rs

@@ -15,7 +15,7 @@ use crate::{
         InsertManyOptions,
         WriteConcern,
     },
-    test::{log_uncaptured, util::EventClient, LOCK},
+    test::{log_uncaptured, util::EventClient, CLIENT_OPTIONS, LOCK},
     Collection,
     Database,
     RUNTIME,
@@ -27,7 +27,10 @@ async fn run_test<F: Future>(
 ) {
     let _guard: RwLockWriteGuard<()> = LOCK.run_exclusively().await;
 
-    let options = ClientOptions::builder().retry_writes(false).build();
+    let options = ClientOptions::builder()
+        .hosts(CLIENT_OPTIONS.hosts.clone())
+        .retry_writes(false)
+        .build();
     let client = EventClient::with_additional_options(Some(options), None, None, None).await;
 
     if !client.is_replica_set() {

src/test/spec/unified_runner/test_runner.rs

@@ -114,12 +114,7 @@ impl TestRunner {
                     options.command_event_handler = Some(observer.clone());
                     options.cmap_event_handler = Some(observer.clone());
                     options.server_api = server_api;
-                    if LOAD_BALANCED_SINGLE_URI
-                        .as_ref()
-                        .map_or(false, |uri| !uri.is_empty())
-                    {
-                        options.test_options_mut().mock_service_id = true;
-                    }
+
                     if TestClient::new().await.is_sharded() {
                         match client.use_multiple_mongoses {
                             Some(true) => {

src/test/util/mod.rs

@@ -73,13 +73,6 @@ impl TestClient {
             options.sdam_event_handler = Some(handler);
         }
 
-        if LOAD_BALANCED_SINGLE_URI
-            .as_ref()
-            .map_or(false, |uri| !uri.is_empty())
-        {
-            options.test_options_mut().mock_service_id = true;
-        }
-
         let client = Client::with_options(options.clone()).unwrap();
 
         // To avoid populating the session pool with leftover implicit sessions, we check out a