RUST-978 Unified test format additions in support of load balanced tests

src/client/mod.rs

@@ -360,4 +360,9 @@ impl Client {
             .map(|stream_address| format!("{}", stream_address))
             .collect()
     }
+
+    #[cfg(test)]
+    pub(crate) async fn sync_workers(&self) {
+        self.inner.topology.sync_workers().await;
+    }
 }

src/cmap/manager.rs

@@ -1,4 +1,6 @@
 use tokio::sync::mpsc;
+#[cfg(test)]
+use tokio::sync::oneshot;
 
 use super::Connection;
 use crate::{bson::oid::ObjectId, error::Error, runtime::AcknowledgedMessage};
@@ -72,6 +74,14 @@ impl PoolManager {
             .sender
             .send(PoolManagementRequest::HandleConnectionSucceeded(conn));
     }
+
+    /// Create a synchronization point for the pool's worker.
+    #[cfg(test)]
+    pub(super) fn sync_worker(&self) -> oneshot::Receiver<()> {
+        let (tx, rx) = oneshot::channel();
+        let _ = self.sender.send(PoolManagementRequest::Sync(tx));
+        rx
+    }
 }
 
 #[derive(Debug)]
@@ -108,6 +118,10 @@ pub(super) enum PoolManagementRequest {
     /// Update the pool after a successful connection, optionally populating the pool
     /// with the successful connection.
     HandleConnectionSucceeded(ConnectionSucceeded),
+
+    /// Synchronize the worker queue state with an external caller, i.e. a test.
+    #[cfg(test)]
+    Sync(oneshot::Sender<()>),
 }
 
 impl PoolManagementRequest {

src/cmap/mod.rs

@@ -12,6 +12,8 @@ mod worker;
 use std::sync::Arc;
 
 use derivative::Derivative;
+#[cfg(test)]
+use tokio::sync::oneshot;
 
 pub use self::conn::ConnectionInfo;
 pub(crate) use self::{
@@ -172,4 +174,9 @@ impl ConnectionPool {
     pub(crate) fn generation(&self) -> PoolGeneration {
         self.generation_subscriber.generation()
     }
+
+    #[cfg(test)]
+    pub(crate) fn sync_worker(&self) -> oneshot::Receiver<()> {
+        self.manager.sync_worker()
+    }
 }

src/cmap/worker.rs

@@ -313,6 +313,10 @@ impl ConnectionPoolWorker {
                 PoolTask::HandleManagementRequest(
                     PoolManagementRequest::HandleConnectionFailed,
                 ) => self.handle_connection_failed(),
+                #[cfg(test)]
+                PoolTask::HandleManagementRequest(PoolManagementRequest::Sync(tx)) => {
+                    let _ = tx.send(());
+                }
                 PoolTask::Maintenance => {
                     self.perform_maintenance();
                 }

src/cursor/common.rs

@@ -8,6 +8,8 @@ use std::{
 use derivative::Derivative;
 use futures_core::{Future, Stream};
 use serde::de::DeserializeOwned;
+#[cfg(test)]
+use tokio::sync::oneshot;
 
 use crate::{
     bson::Document,
@@ -303,13 +305,18 @@ pub(super) fn kill_cursor(
     ns: &Namespace,
     cursor_id: i64,
     pinned_conn: PinnedConnection,
+    #[cfg(test)] kill_watcher: Option<oneshot::Sender<()>>,
 ) {
     let coll = client
         .database(ns.db.as_str())
         .collection::<Document>(ns.coll.as_str());
     RUNTIME.execute(async move {
         if !pinned_conn.is_invalid() {
             let _ = coll.kill_cursor(cursor_id, pinned_conn.handle()).await;
+            #[cfg(test)]
+            if let Some(tx) = kill_watcher {
+                let _ = tx.send(());
+            }
         }
         pinned_conn.unpin().await;
     });

src/cursor/mod.rs

@@ -8,6 +8,8 @@ use std::{
 
 use futures_core::{future::BoxFuture, Stream};
 use serde::de::DeserializeOwned;
+#[cfg(test)]
+use tokio::sync::oneshot;
 
 use crate::{
     cmap::conn::PinnedConnectionHandle,
@@ -84,6 +86,8 @@ where
 {
     client: Client,
     wrapped_cursor: ImplicitSessionCursor<T>,
+    #[cfg(test)]
+    kill_watcher: Option<oneshot::Sender<()>>,
     _phantom: std::marker::PhantomData<T>,
 }
 
@@ -107,9 +111,24 @@ where
                 PinnedConnection::new(pin),
                 provider,
             ),
+            #[cfg(test)]
+            kill_watcher: None,
             _phantom: Default::default(),
         }
     }
+
+    /// Some tests need to be able to observe the events generated by `killCommand` execution;
+    /// however, because that happens asynchronously on `drop`, the test runner can conclude before
+    /// the event is published.  To fix that, tests can set a "kill watcher" on cursors - a
+    /// one-shot channel with a `()` value pushed after `killCommand` is run that the test can wait
+    /// on.
+    #[cfg(test)]
+    pub(crate) fn set_kill_watcher(&mut self, tx: oneshot::Sender<()>) {
+        if self.kill_watcher.is_some() {
+            panic!("cursor already has a kill_watcher");
+        }
+        self.kill_watcher = Some(tx);
+    }
 }
 
 impl<T> Stream for Cursor<T>
@@ -137,6 +156,8 @@ where
             self.wrapped_cursor.namespace(),
             self.wrapped_cursor.id(),
             self.wrapped_cursor.pinned_connection().replicate(),
+            #[cfg(test)]
+            self.kill_watcher.take(),
         );
     }
 }

src/cursor/session.rs

@@ -7,6 +7,8 @@ use std::{
 use futures_core::{future::BoxFuture, Stream};
 use futures_util::StreamExt;
 use serde::de::DeserializeOwned;
+#[cfg(test)]
+use tokio::sync::oneshot;
 
 use super::common::{
     kill_cursor,
@@ -64,6 +66,8 @@ where
     info: CursorInformation,
     buffer: VecDeque<T>,
     pinned_connection: PinnedConnection,
+    #[cfg(test)]
+    kill_watcher: Option<oneshot::Sender<()>>,
 }
 
 impl<T> SessionCursor<T>
@@ -83,6 +87,8 @@ where
             info: spec.info,
             buffer: spec.initial_buffer,
             pinned_connection: PinnedConnection::new(pinned),
+            #[cfg(test)]
+            kill_watcher: None,
         }
     }
 
@@ -177,6 +183,19 @@ where
     pub async fn next(&mut self, session: &mut ClientSession) -> Option<Result<T>> {
         self.stream(session).next().await
     }
+
+    /// Some tests need to be able to observe the events generated by `killCommand` execution;
+    /// however, because that happens asynchronously on `drop`, the test runner can conclude before
+    /// the event is published.  To fix that, tests can set a "kill watcher" on cursors - a
+    /// one-shot channel with a `()` value pushed after `killCommand` is run that the test can wait
+    /// on.
+    #[cfg(test)]
+    pub(crate) fn set_kill_watcher(&mut self, tx: oneshot::Sender<()>) {
+        if self.kill_watcher.is_some() {
+            panic!("cursor already has a kill_watcher");
+        }
+        self.kill_watcher = Some(tx);
+    }
 }
 
 impl<T> Drop for SessionCursor<T>
@@ -193,6 +212,8 @@ where
             &self.info.ns,
             self.info.id,
             self.pinned_connection.replicate(),
+            #[cfg(test)]
+            self.kill_watcher.take(),
         );
     }
 }

src/operation/count/mod.rs

@@ -47,8 +47,10 @@ impl Operation for Count {
     const NAME: &'static str = "count";
 
     fn build(&mut self, description: &StreamDescription) -> Result<Command> {
+        let mut name = Self::NAME.to_string();
         let mut body = match description.max_wire_version {
             Some(v) if v >= SERVER_4_9_0_WIRE_VERSION => {
+                name = "aggregate".to_string();
                 doc! {
                     "aggregate": self.ns.coll.clone(),
                     "pipeline": [
@@ -74,11 +76,7 @@ impl Operation for Count {
 
         append_options(&mut body, self.options.as_ref())?;
 
-        Ok(Command::new(
-            Self::NAME.to_string(),
-            self.ns.db.clone(),
-            body,
-        ))
+        Ok(Command::new(name, self.ns.db.clone(), body))
     }
 
     fn handle_response(

src/sdam/state/mod.rs

@@ -9,6 +9,8 @@ use std::{
     },
 };
 
+#[cfg(test)]
+use futures_util::stream::{FuturesUnordered, StreamExt};
 use tokio::sync::{RwLock, RwLockWriteGuard};
 
 use self::server::Server;
@@ -532,6 +534,11 @@ impl Topology {
             .map(|(addr, server)| (addr.clone(), Arc::downgrade(server)))
             .collect()
     }
+
+    #[cfg(test)]
+    pub(crate) async fn sync_workers(&self) {
+        self.state.read().await.sync_workers().await;
+    }
 }
 
 impl WeakTopology {
@@ -680,6 +687,16 @@ impl TopologyState {
 
         self.servers.retain(|host, _| hosts.contains(host));
     }
+
+    #[cfg(test)]
+    async fn sync_workers(&self) {
+        let rxen: FuturesUnordered<_> = self
+            .servers
+            .values()
+            .map(|v| v.pool.sync_worker())
+            .collect();
+        let _: Vec<_> = rxen.collect().await;
+    }
 }
 
 /// Enum describing a point in time during an operation's execution relative to when the MongoDB

src/test/mod.rs

@@ -66,6 +66,10 @@ lazy_static! {
     };
     pub(crate) static ref SERVERLESS: bool =
         matches!(std::env::var("SERVERLESS"), Ok(s) if s == "serverless");
+    pub(crate) static ref LOAD_BALANCED_SINGLE_URI: Option<String> =
+        std::env::var("MONGODB_LOAD_BALANCED_SINGLE_URI").ok();
+    pub(crate) static ref LOAD_BALANCED_MULTIPLE_URI: Option<String> =
+        std::env::var("MONGODB_LOAD_BALANCED_MULTIPLE_URI").ok();
 }
 
 fn get_default_uri() -> String {

src/test/spec/command_monitoring/operation.rs

@@ -114,6 +114,8 @@ impl TestOperation for DeleteOne {
 #[serde(deny_unknown_fields)]
 pub(super) struct Find {
     filter: Option<Document>,
+    // `FindOptions` cannot be embedded directly because serde doesn't support combining `flatten`
+    // and `deny_unknown_fields`, so its fields are replicated here.
     #[serde(default)]
     sort: Option<Document>,
     #[serde(default)]
@@ -149,6 +151,8 @@ impl TestOperation for Find {
 
     fn execute(&self, collection: Collection<Document>) -> BoxFuture<Result<()>> {
         async move {
+            // `FindOptions` is constructed without the use of `..Default::default()` to enforce at
+            // compile-time that any new fields added there need to be considered here.
             let options = FindOptions {
                 sort: self.sort.clone(),
                 skip: self.skip,
@@ -161,7 +165,16 @@ impl TestOperation for Find {
                 max: self.max.clone(),
                 return_key: self.return_key,
                 show_record_id: self.show_record_id,
-                ..Default::default()
+                allow_disk_use: None,
+                allow_partial_results: None,
+                cursor_type: None,
+                max_await_time: None,
+                max_scan: None,
+                no_cursor_timeout: None,
+                projection: None,
+                read_concern: None,
+                selection_criteria: None,
+                collation: None,
             };
 
             let mut cursor = collection.find(self.filter.clone(), options).await?;

src/test/spec/connection_stepdown.rs

@@ -105,7 +105,7 @@ async fn get_more() {
         }
 
         RUNTIME.delay_for(Duration::from_millis(250)).await;
-        assert!(client.get_pool_cleared_events().is_empty());
+        assert_eq!(client.count_pool_cleared_events(), 0);
     }
 
     run_test(function_name!(), get_more_test).await;
@@ -155,7 +155,7 @@ async fn not_master_keep_pool() {
             .expect("insert should have succeeded");
 
         RUNTIME.delay_for(Duration::from_millis(250)).await;
-        assert!(client.get_pool_cleared_events().is_empty());
+        assert_eq!(client.count_pool_cleared_events(), 0);
     }
 
     run_test(function_name!(), not_master_keep_pool_test).await;
@@ -201,7 +201,7 @@ async fn not_master_reset_pool() {
         );
 
         RUNTIME.delay_for(Duration::from_millis(250)).await;
-        assert!(client.get_pool_cleared_events().len() == 1);
+        assert_eq!(client.count_pool_cleared_events(), 1);
 
         coll.insert_one(doc! { "test": 1 }, None)
             .await
@@ -250,7 +250,7 @@ async fn shutdown_in_progress() {
         );
 
         RUNTIME.delay_for(Duration::from_millis(250)).await;
-        assert!(client.get_pool_cleared_events().len() == 1);
+        assert_eq!(client.count_pool_cleared_events(), 1);
 
         coll.insert_one(doc! { "test": 1 }, None)
             .await
@@ -299,7 +299,7 @@ async fn interrupted_at_shutdown() {
         );
 
         RUNTIME.delay_for(Duration::from_millis(250)).await;
-        assert!(client.get_pool_cleared_events().len() == 1);
+        assert_eq!(client.count_pool_cleared_events(), 1);
 
         coll.insert_one(doc! { "test": 1 }, None)
             .await

src/test/spec/json/unified-test-format/invalid/expectedCmapEvent-connectionCheckOutFailedEvent-reason-type.json

@@ -0,0 +1,23 @@
+{
+  "description": "expectedCmapEvent-connectionCheckOutFailedEvent-reason-type",
+  "schemaVersion": "1.3",
+  "tests": [
+    {
+      "description": "foo",
+      "operations": [],
+      "expectEvents": [
+        {
+          "client": "client0",
+          "eventType": "cmap",
+          "events": [
+            {
+              "connectionCheckOutFailedEvent": {
+                "reason": 10
+              }
+            }
+          ]
+        }
+      ]
+    }
+  ]
+}