RUST-1149 Prose tests for change streams. (#561)

Author: abr-egn

src/change_stream/event.rs

@@ -22,7 +22,7 @@ use serde::{Deserialize, Serialize};
 /// See the documentation
 /// [here](https://docs.mongodb.com/manual/changeStreams/#change-stream-resume-token) for more
 /// information on resume tokens.
-#[derive(Clone, Debug, Deserialize, Serialize)]
+#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
 pub struct ResumeToken(pub(crate) RawBson);
 
 impl ResumeToken {
@@ -43,11 +43,16 @@ impl ResumeToken {
     pub(crate) fn from_raw(doc: Option<RawDocumentBuf>) -> Option<ResumeToken> {
         doc.map(|doc| ResumeToken(RawBson::Document(doc)))
     }
+
+    #[cfg(test)]
+    pub fn parsed(self) -> std::result::Result<Bson, bson::raw::Error> {
+        self.0.try_into()
+    }
 }
 
 /// A `ChangeStreamEvent` represents a
 /// [change event](https://docs.mongodb.com/manual/reference/change-events/) in the associated change stream.
-#[derive(Debug, Deserialize)]
+#[derive(Debug, Deserialize, PartialEq)]
 #[serde(rename_all = "camelCase")]
 #[non_exhaustive]
 pub struct ChangeStreamEvent<T> {
@@ -99,7 +104,7 @@ pub struct ChangeStreamEvent<T> {
 }
 
 /// Describes which fields have been updated or removed from a document.
-#[derive(Debug, Deserialize)]
+#[derive(Debug, Deserialize, PartialEq)]
 #[serde(rename_all = "camelCase")]
 #[non_exhaustive]
 pub struct UpdateDescription {
@@ -112,7 +117,7 @@ pub struct UpdateDescription {
 }
 
 /// The operation type represented in a given change notification.
-#[derive(Debug, Deserialize, Clone, PartialEq)]
+#[derive(Debug, Deserialize, Clone, PartialEq, Eq)]
 #[serde(rename_all = "camelCase")]
 #[non_exhaustive]
 pub enum OperationType {
@@ -142,13 +147,12 @@ pub enum OperationType {
 }
 
 /// Identifies the collection or database on which an event occurred.
-#[derive(Deserialize, Debug)]
-#[serde(untagged)]
+#[derive(Deserialize, Debug, PartialEq, Eq)]
 #[non_exhaustive]
-pub enum ChangeStreamEventSource {
-    /// The [`Namespace`] containing the database and collection in which the change occurred.
-    Namespace(Namespace),
+pub struct ChangeStreamEventSource {
+    /// The name of the database in which the change occurred.
+    pub db: String,
 
-    /// Contains the name of the database in which the change happened.
-    Database(String),
+    /// The name of the collection in which the change occurred.
+    pub coll: Option<String>,
 }

src/change_stream/mod.rs

@@ -21,7 +21,7 @@ use crate::{
         options::ChangeStreamOptions,
     },
     cursor::{stream_poll_next, BatchValue, CursorStream, NextInBatchFuture},
-    error::{Error, Result},
+    error::{Error, ErrorKind, Result},
     operation::AggregateTarget,
     options::AggregateOptions,
     selection_criteria::{ReadPreference, SelectionCriteria},
@@ -220,10 +220,14 @@ fn get_resume_token(
 ) -> Result<Option<ResumeToken>> {
     Ok(match batch_value {
         BatchValue::Some { doc, is_last } => {
+            let doc_token = match doc.get("_id")? {
+                Some(val) => ResumeToken(val.to_raw_bson()),
+                None => return Err(ErrorKind::MissingResumeToken.into()),
+            };
             if *is_last && batch_token.is_some() {
                 batch_token.cloned()
             } else {
-                doc.get("_id")?.map(|val| ResumeToken(val.to_raw_bson()))
+                Some(doc_token)
             }
         }
         BatchValue::Empty => batch_token.cloned(),
@@ -236,50 +240,56 @@ where
     T: DeserializeOwned + Unpin + Send + Sync,
 {
     fn poll_next_in_batch(&mut self, cx: &mut Context<'_>) -> Poll<Result<BatchValue>> {
-        if let Some(mut pending) = self.pending_resume.take() {
-            match Pin::new(&mut pending).poll(cx) {
-                Poll::Pending => {
-                    self.pending_resume = Some(pending);
-                    return Poll::Pending;
+        loop {
+            if let Some(mut pending) = self.pending_resume.take() {
+                match Pin::new(&mut pending).poll(cx) {
+                    Poll::Pending => {
+                        self.pending_resume = Some(pending);
+                        return Poll::Pending;
+                    }
+                    Poll::Ready(Ok(new_stream)) => {
+                        // Ensure that the old cursor is killed on the server selected for the new
+                        // one.
+                        self.cursor
+                            .set_drop_address(new_stream.cursor.address().clone());
+                        self.cursor = new_stream.cursor;
+                        self.args = new_stream.args;
+                        continue;
+                    }
+                    Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
                 }
-                Poll::Ready(Ok(new_stream)) => {
-                    // Ensure that the old cursor is killed on the server selected for the new one.
-                    self.cursor
-                        .set_drop_address(new_stream.cursor.address().clone());
-                    self.cursor = new_stream.cursor;
-                    self.args = new_stream.args;
-                    return Poll::Pending;
-                }
-                Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
             }
-        }
-        let out = self.cursor.poll_next_in_batch(cx);
-        match &out {
-            Poll::Ready(Ok(bv)) => {
-                if let Some(token) = get_resume_token(bv, self.cursor.post_batch_resume_token())? {
-                    self.data.resume_token = Some(token);
+            let out = self.cursor.poll_next_in_batch(cx);
+            match &out {
+                Poll::Ready(Ok(bv)) => {
+                    if let Some(token) =
+                        get_resume_token(bv, self.cursor.post_batch_resume_token())?
+                    {
+                        self.data.resume_token = Some(token);
+                    }
+                    if matches!(bv, BatchValue::Some { .. }) {
+                        self.data.document_returned = true;
+                    }
                 }
-                if matches!(bv, BatchValue::Some { .. }) {
-                    self.data.document_returned = true;
+                Poll::Ready(Err(e)) if e.is_resumable() && !self.data.resume_attempted => {
+                    self.data.resume_attempted = true;
+                    let client = self.cursor.client().clone();
+                    let args = self.args.clone();
+                    let mut data = self.data.take();
+                    data.implicit_session = self.cursor.take_implicit_session();
+                    self.pending_resume = Some(Box::pin(async move {
+                        let new_stream: Result<ChangeStream<ChangeStreamEvent<()>>> = client
+                            .execute_watch(args.pipeline, args.options, args.target, Some(data))
+                            .await;
+                        new_stream.map(|cs| cs.with_type::<T>())
+                    }));
+                    // Iterate the loop so the new future gets polled and can register wakers.
+                    continue;
                 }
+                _ => {}
             }
-            Poll::Ready(Err(e)) if e.is_resumable() && !self.data.resume_attempted => {
-                self.data.resume_attempted = true;
-                let client = self.cursor.client().clone();
-                let args = self.args.clone();
-                let mut data = self.data.take();
-                data.implicit_session = self.cursor.take_implicit_session();
-                self.pending_resume = Some(Box::pin(async move {
-                    let new_stream: Result<ChangeStream<ChangeStreamEvent<()>>> = client
-                        .execute_watch(args.pipeline, args.options, args.target, Some(data))
-                        .await;
-                    new_stream.map(|cs| cs.with_type::<T>())
-                }));
-                return Poll::Pending;
-            }
-            _ => {}
+            return out;
         }
-        out
     }
 }
 

src/cmap/test/integration.rs

@@ -180,7 +180,7 @@ async fn connection_error_during_establishment() {
     client_options.repl_set_name = None;
 
     let client = TestClient::with_options(Some(client_options.clone())).await;
-    if !client.supports_fail_command().await {
+    if !client.supports_fail_command() {
         println!(
             "skipping {} due to failCommand not being supported",
             function_name!()
@@ -235,7 +235,7 @@ async fn connection_error_during_operation() {
     options.max_pool_size = Some(1);
 
     let client = TestClient::with_options(options.into()).await;
-    if !client.supports_fail_command().await {
+    if !client.supports_fail_command() {
         println!(
             "skipping {} due to failCommand not being supported",
             function_name!()

src/coll/mod.rs

@@ -1354,7 +1354,7 @@ where
 }
 
 /// A struct modeling the canonical name for a collection in MongoDB.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub struct Namespace {
     /// The name of the database associated with this namespace.
     pub db: String,

src/error.rs

@@ -470,6 +470,10 @@ pub enum ErrorKind {
     #[error("The server does not support a database operation: {message}")]
     #[non_exhaustive]
     IncompatibleServer { message: String },
+
+    /// No resume token was present in a change stream document.
+    #[error("Cannot provide resume functionality when the resume token is missing")]
+    MissingResumeToken,
 }
 
 impl ErrorKind {

src/sdam/description/topology/test/sdam.rs

@@ -615,7 +615,7 @@ async fn heartbeat_events() {
         .await
         .expect("should see server heartbeat succeeded event");
 
-    if !client.supports_fail_command().await {
+    if !client.supports_fail_command() {
         return;
     }
 

src/sdam/test.rs

@@ -42,7 +42,7 @@ async fn min_heartbeat_frequency() {
 
     let setup_client = TestClient::with_options(Some(setup_client_options.clone())).await;
 
-    if !setup_client.supports_fail_command().await {
+    if !setup_client.supports_fail_command() {
         println!("skipping min_heartbeat_frequency test due to server not supporting fail points");
         return;
     }