use lb uri for multiple mongoses in more places

src/test/mod.rs

@@ -44,20 +44,7 @@ use std::{fs::read_to_string, str::FromStr};
 const MAX_POOL_SIZE: u32 = 100;
 
 lazy_static! {
-    pub(crate) static ref CLIENT_OPTIONS: ClientOptions = {
-        let uri = DEFAULT_URI.clone();
-        let mut options = ClientOptions::parse_without_srv_resolution(&uri).unwrap();
-        options.max_pool_size = Some(MAX_POOL_SIZE);
-        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;
-        }
-
-        options
-    };
+    pub(crate) static ref CLIENT_OPTIONS: ClientOptions = client_options_for_uri(&DEFAULT_URI);
     pub(crate) static ref LOCK: TestLock = TestLock::new();
     pub(crate) static ref DEFAULT_URI: String = get_default_uri();
     pub(crate) static ref SERVER_API: Option<ServerApi> = match std::env::var("MONGODB_API_VERSION")
@@ -77,6 +64,20 @@ lazy_static! {
         std::env::var("MULTI_MONGOS_LB_URI").ok();
 }
 
+pub(crate) fn client_options_for_uri(uri: &str) -> ClientOptions {
+    let mut options = ClientOptions::parse_without_srv_resolution(uri).unwrap();
+    options.max_pool_size = Some(MAX_POOL_SIZE);
+    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;
+    }
+
+    options
+}
+
 fn get_default_uri() -> String {
     if let Some(uri) = LOAD_BALANCED_SINGLE_URI.clone() {
         if !uri.is_empty() {

src/test/util/event.rs

@@ -47,7 +47,7 @@ use crate::{
         },
     },
     options::ClientOptions,
-    test::{spec::ExpectedEventType, CLIENT_OPTIONS, LOCK},
+    test::{spec::ExpectedEventType, LOCK},
     RUNTIME,
 };
 
@@ -448,20 +448,11 @@ impl EventClient {
         use_multiple_mongoses: Option<bool>,
         event_handler: impl Into<Option<EventHandler>>,
     ) -> Self {
-        let mut options = match options.into() {
-            Some(mut options) => {
-                options.merge(CLIENT_OPTIONS.clone());
-                options
-            }
-            None => CLIENT_OPTIONS.clone(),
-        };
+        let mut options = TestClient::options_for_multiple_mongoses(options.into(), use_multiple_mongoses.unwrap_or(false)).await;
         options
             .test_options
             .get_or_insert_with(Default::default)
             .heartbeat_freq = heartbeat_freq;
-        if TestClient::new().await.is_sharded() && use_multiple_mongoses != Some(true) {
-            options.hosts = options.hosts.iter().cloned().take(1).collect();
-        }
         EventClient::with_options_and_handler(options, event_handler).await
     }
 

src/test/util/mod.rs

@@ -25,7 +25,7 @@ use crate::{
     error::{CommandError, ErrorKind, Result},
     operation::RunCommand,
     options::{AuthMechanism, ClientOptions, CollectionOptions, CreateCollectionOptions},
-    test::{Topology, LOAD_BALANCED_SINGLE_URI},
+    test::{Topology, LOAD_BALANCED_SINGLE_URI, LOAD_BALANCED_MULTIPLE_URI, client_options_for_uri},
     Client,
     Collection,
 };
@@ -127,16 +127,7 @@ impl TestClient {
         options: Option<ClientOptions>,
         use_multiple_mongoses: bool,
     ) -> Self {
-        let mut options = match options {
-            Some(mut options) => {
-                options.merge(CLIENT_OPTIONS.clone());
-                options
-            }
-            None => CLIENT_OPTIONS.clone(),
-        };
-        if !use_multiple_mongoses && Self::new().await.is_sharded() {
-            options.hosts = options.hosts.iter().cloned().take(1).collect();
-        }
+        let options = Self::options_for_multiple_mongoses(options, use_multiple_mongoses).await;
         Self::with_options(Some(options)).await
     }
 
@@ -364,6 +355,35 @@ impl TestClient {
         }
         .to_string()
     }
+
+    pub async fn options_for_multiple_mongoses(options: Option<ClientOptions>, use_multiple_mongoses: bool) -> ClientOptions {
+        let is_load_balanced = options
+            .as_ref()
+            .and_then(|o| o.load_balanced)
+            .or_else(|| CLIENT_OPTIONS.load_balanced)
+            .unwrap_or(false);
+        let default_options = if is_load_balanced {
+            let uri = if use_multiple_mongoses {
+                LOAD_BALANCED_MULTIPLE_URI.as_ref().expect("MULTI_MONGOS_LB_URI is required")
+            } else {
+                LOAD_BALANCED_SINGLE_URI.as_ref().expect("SINGLE_MONGOS_LB_URI is required")
+            };
+            client_options_for_uri(uri)
+        } else {
+            CLIENT_OPTIONS.clone()
+        };
+        let mut options = match options {
+            Some(mut options) => {
+                options.merge(default_options);
+                options
+            }
+            None => default_options,
+        };
+        if Self::new().await.is_sharded() && !use_multiple_mongoses {
+            options.hosts = options.hosts.iter().cloned().take(1).collect();
+        }
+        options
+    }
 }
 
 pub async fn drop_collection<T>(coll: &Collection<T>)