RUST-1191 Ensure change stream is in proper state in mid_batch_resume_token test (#689)

Author: patrickfreed

src/change_stream/mod.rs

@@ -3,12 +3,16 @@ pub mod event;
 pub(crate) mod options;
 pub mod session;
 
+#[cfg(test)]
+use std::collections::VecDeque;
 use std::{
     future::Future,
     pin::Pin,
     task::{Context, Poll},
 };
 
+#[cfg(test)]
+use bson::RawDocumentBuf;
 use bson::{Document, Timestamp};
 use derivative::Derivative;
 use futures_core::{future::BoxFuture, Stream};
@@ -167,6 +171,11 @@ where
     pub(crate) fn set_kill_watcher(&mut self, tx: oneshot::Sender<()>) {
         self.cursor.set_kill_watcher(tx);
     }
+
+    #[cfg(test)]
+    pub(crate) fn current_batch(&self) -> &VecDeque<RawDocumentBuf> {
+        self.cursor.current_batch()
+    }
 }
 
 /// Arguments passed to a `watch` method, captured to allow resume.

src/cursor/common.rs

@@ -87,6 +87,11 @@ where
         self.state().buffer.current()
     }
 
+    #[cfg(test)]
+    pub(super) fn current_batch(&self) -> &VecDeque<RawDocumentBuf> {
+        self.state().buffer.as_ref()
+    }
+
     fn state_mut(&mut self) -> &mut CursorState {
         self.state.as_mut().unwrap()
     }
@@ -531,3 +536,9 @@ impl CursorBuffer {
         self.docs.front().map(|d| d.as_ref())
     }
 }
+
+impl AsRef<VecDeque<RawDocumentBuf>> for CursorBuffer {
+    fn as_ref(&self) -> &VecDeque<RawDocumentBuf> {
+        &self.docs
+    }
+}

src/cursor/mod.rs

@@ -1,12 +1,17 @@
 mod common;
 pub(crate) mod session;
 
+#[cfg(test)]
+use std::collections::VecDeque;
 use std::{
     pin::Pin,
     task::{Context, Poll},
 };
 
 use bson::RawDocument;
+
+#[cfg(test)]
+use bson::RawDocumentBuf;
 use futures_core::{future::BoxFuture, Stream};
 use serde::{de::DeserializeOwned, Deserialize};
 #[cfg(test)]
@@ -269,6 +274,11 @@ impl<T> Cursor<T> {
         );
         self.kill_watcher = Some(tx);
     }
+
+    #[cfg(test)]
+    pub(crate) fn current_batch(&self) -> &VecDeque<RawDocumentBuf> {
+        self.wrapped_cursor.as_ref().unwrap().current_batch()
+    }
 }
 
 impl<T> CursorStream for Cursor<T>

src/test/change_stream.rs

@@ -455,11 +455,28 @@ async fn batch_mid_resume_token() -> Result<()> {
         None => return Ok(()),
     };
 
-    coll.insert_many((0..2).map(|i| doc! { "_id": i as i32 }), None)
-        .await?;
+    // This loop gets the stream to a point where it has been iterated up to but not including
+    // the last event in its batch.
+    let mut event_id = None;
+    loop {
+        match stream.next_if_any().await? {
+            Some(event) => {
+                event_id = Some(event.id);
+            }
+            // If we're out of events, make some more.
+            None => {
+                coll.insert_many((0..3).map(|_| doc! {}), None).await?;
+            }
+        };
+
+        // if after iterating the stream last time there's one document left,
+        // then we're done here and can continue to the assertions.
+        if stream.current_batch().len() == 1 {
+            break;
+        }
+    }
 
-    let mid_id = stream.next().await.transpose()?.unwrap().id;
-    assert_eq!(stream.resume_token(), Some(mid_id));
+    assert_eq!(stream.resume_token().unwrap(), event_id.unwrap());
 
     Ok(())
 }