RUST-1855 Propagate cluster time for unified tests (#1053)

Author: abr-egn

src/test/coll.rs

@@ -32,11 +32,7 @@ use crate::{
         WriteConcern,
     },
     results::DeleteResult,
-    test::{
-        get_client_options,
-        log_uncaptured,
-        util::{drop_collection, TestClient},
-    },
+    test::{get_client_options, log_uncaptured, util::TestClient},
     Collection,
     IndexModel,
 };
@@ -209,7 +205,7 @@ async fn aggregate_out() {
     let db = client.database(function_name!());
     let coll = db.collection(function_name!());
 
-    drop_collection(&coll).await;
+    coll.drop().await.unwrap();
 
     let result = coll
         .insert_many((0i32..5).map(|n| doc! { "x": n }).collect::<Vec<_>>())
@@ -226,7 +222,7 @@ async fn aggregate_out() {
         },
         doc! {"$out": out_coll.name()},
     ];
-    drop_collection(&out_coll).await;
+    out_coll.drop().await.unwrap();
 
     coll.aggregate(pipeline.clone()).await.unwrap();
     assert!(db
@@ -235,7 +231,7 @@ async fn aggregate_out() {
         .unwrap()
         .into_iter()
         .any(|name| name.as_str() == out_coll.name()));
-    drop_collection(&out_coll).await;
+    out_coll.drop().await.unwrap();
 
     // check that even with a batch size of 0, a new collection is created.
     coll.aggregate(pipeline).batch_size(0).await.unwrap();
@@ -262,7 +258,7 @@ async fn kill_cursors_on_drop() {
     let db = client.database(function_name!());
     let coll = db.collection(function_name!());
 
-    drop_collection(&coll).await;
+    coll.drop().await.unwrap();
 
     coll.insert_many(vec![doc! { "x": 1 }, doc! { "x": 2 }])
         .await
@@ -295,7 +291,7 @@ async fn no_kill_cursors_on_exhausted() {
     let db = client.database(function_name!());
     let coll = db.collection(function_name!());
 
-    drop_collection(&coll).await;
+    coll.drop().await.unwrap();
 
     coll.insert_many(vec![doc! { "x": 1 }, doc! { "x": 2 }])
         .await

src/test/spec/unified_runner/test_runner.rs

@@ -7,15 +7,9 @@ use tokio::sync::{mpsc, RwLock};
 use crate::{
     bson::{doc, Document},
     client::options::ClientOptions,
-    concern::{Acknowledgment, WriteConcern},
+    concern::WriteConcern,
     gridfs::GridFsBucket,
-    options::{
-        CollectionOptions,
-        CreateCollectionOptions,
-        ReadConcern,
-        ReadPreference,
-        SelectionCriteria,
-    },
+    options::{CollectionOptions, ReadConcern, ReadPreference, SelectionCriteria},
     runtime,
     sdam::{TopologyDescription, MIN_HEARTBEAT_FREQUENCY},
     test::{
@@ -35,6 +29,8 @@ use crate::{
         SERVERLESS,
         SERVER_API,
     },
+    ClientSession,
+    ClusterTime,
     Collection,
     Database,
 };
@@ -82,6 +78,7 @@ pub(crate) struct TestRunner {
     pub(crate) internal_client: TestClient,
     pub(crate) entities: Arc<RwLock<EntityMap>>,
     pub(crate) fail_point_guards: Arc<RwLock<Vec<FailPointGuard>>>,
+    pub(crate) cluster_time: Arc<RwLock<Option<ClusterTime>>>,
 }
 
 impl TestRunner {
@@ -90,6 +87,7 @@ impl TestRunner {
             internal_client: TestClient::new().await,
             entities: Default::default(),
             fail_point_guards: Default::default(),
+            cluster_time: Default::default(),
         }
     }
 
@@ -99,6 +97,7 @@ impl TestRunner {
             internal_client: TestClient::with_options(Some(options)).await,
             entities: Arc::new(RwLock::new(EntityMap::new())),
             fail_point_guards: Arc::new(RwLock::new(Vec::new())),
+            cluster_time: Default::default(),
         }
     }
 
@@ -205,9 +204,11 @@ impl TestRunner {
             log_uncaptured(format!("Executing {:?}", &test_case.description));
 
             if let Some(ref initial_data) = test_file.initial_data {
+                let mut session = self.internal_client.start_session().await.unwrap();
                 for data in initial_data {
-                    self.insert_initial_data(data).await;
+                    self.insert_initial_data(data, &mut session).await;
                 }
+                *self.cluster_time.write().await = session.cluster_time().cloned();
             }
 
             self.entities.write().await.clear();
@@ -370,33 +371,37 @@ impl TestRunner {
         }
     }
 
-    pub(crate) async fn insert_initial_data(&self, data: &CollectionData) {
-        let write_concern = WriteConcern::builder().w(Acknowledgment::Majority).build();
-
+    pub(crate) async fn insert_initial_data(
+        &self,
+        data: &CollectionData,
+        session: &mut ClientSession,
+    ) {
         if !data.documents.is_empty() {
             let collection_options = CollectionOptions::builder()
-                .write_concern(write_concern)
+                .write_concern(WriteConcern::majority())
                 .build();
+            let coll = self.internal_client.get_coll_with_options(
+                &data.database_name,
+                &data.collection_name,
+                collection_options,
+            );
+            coll.drop().session(&mut *session).await.unwrap();
+            coll.insert_many(data.documents.clone())
+                .session(session)
+                .await
+                .unwrap();
+        } else {
             let coll = self
                 .internal_client
-                .init_db_and_coll_with_options(
-                    &data.database_name,
-                    &data.collection_name,
-                    collection_options,
-                )
-                .await;
-            coll.insert_many(data.documents.clone()).await.unwrap();
-        } else {
-            let collection_options = CreateCollectionOptions::builder()
-                .write_concern(write_concern)
-                .build();
+                .get_coll(&data.database_name, &data.collection_name);
+            coll.drop().session(&mut *session).await.unwrap();
             self.internal_client
-                .create_fresh_collection(
-                    &data.database_name,
-                    &data.collection_name,
-                    collection_options,
-                )
-                .await;
+                .database(&data.database_name)
+                .create_collection(&data.collection_name)
+                .session(&mut *session)
+                .write_concern(WriteConcern::majority())
+                .await
+                .unwrap();
         }
     }
 
@@ -526,11 +531,14 @@ impl TestRunner {
                 TestFileEntity::Session(session) => {
                     let id = session.id.clone();
                     let client = self.get_client(&session.client).await;
-                    let client_session = client
+                    let mut client_session = client
                         .start_session()
                         .with_options(session.session_options.clone())
                         .await
                         .unwrap();
+                    if let Some(time) = &*self.cluster_time.read().await {
+                        client_session.advance_cluster_time(time);
+                    }
                     (id, Entity::Session(SessionEntity::new(client_session)))
                 }
                 TestFileEntity::Bucket(bucket) => {

src/test/util.rs

@@ -37,7 +37,7 @@ use std::{fmt::Debug, time::Duration};
 
 use super::get_client_options;
 use crate::{
-    error::{CommandError, ErrorKind, Result},
+    error::Result,
     options::{AuthMechanism, ClientOptions, CollectionOptions, CreateCollectionOptions},
     test::{
         update_options_for_testing,
@@ -257,7 +257,7 @@ impl TestClient {
         coll_name: &str,
     ) -> Collection<Document> {
         let coll = self.get_coll(db_name, coll_name);
-        drop_collection(&coll).await;
+        coll.drop().await.unwrap();
         coll
     }
 
@@ -270,7 +270,7 @@ impl TestClient {
         T: Serialize + DeserializeOwned + Unpin + Debug + Send + Sync,
     {
         let coll = self.database(db_name).collection(coll_name);
-        drop_collection(&coll).await;
+        coll.drop().await.unwrap();
         coll
     }
 
@@ -284,17 +284,6 @@ impl TestClient {
             .collection_with_options(coll_name, options)
     }
 
-    pub(crate) async fn init_db_and_coll_with_options(
-        &self,
-        db_name: &str,
-        coll_name: &str,
-        options: CollectionOptions,
-    ) -> Collection<Document> {
-        let coll = self.get_coll_with_options(db_name, coll_name, options);
-        drop_collection(&coll).await;
-        coll
-    }
-
     pub(crate) async fn create_fresh_collection(
         &self,
         db_name: &str,
@@ -410,7 +399,7 @@ impl TestClient {
 
     pub(crate) async fn drop_collection(&self, db_name: &str, coll_name: &str) {
         let coll = self.get_coll(db_name, coll_name);
-        drop_collection(&coll).await;
+        coll.drop().await.unwrap();
     }
 
     /// Returns the `Topology' that can be determined without a server query, i.e. all except
@@ -490,18 +479,6 @@ impl TestClient {
     }
 }
 
-pub(crate) async fn drop_collection<T>(coll: &Collection<T>)
-where
-    T: Serialize + DeserializeOwned + Unpin + Debug + Send + Sync,
-{
-    match coll.drop().await.map_err(|e| *e.kind) {
-        Err(ErrorKind::Command(CommandError { code: 26, .. })) | Ok(_) => {}
-        e @ Err(_) => {
-            e.unwrap();
-        }
-    };
-}
-
 pub(crate) fn get_default_name(description: &str) -> String {
     let mut db_name = description.replace('$', "%").replace([' ', '.'], "_");
     // database names must have fewer than 38 characters