RUST-1757 Fix final cursor batch handling (#951)

mongodb · Sep 15, 2023 · 85de932 · 85de932
1 parent 1e0dd95
commit 85de932
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 6 deletions.
diff --git a/src/cursor/common.rs b/src/cursor/common.rs
@@ -135,12 +135,14 @@ where
         let result = self.provider.execute(spec, client, pin).await;
         self.handle_get_more_result(result)?;
 
-        if self.is_exhausted() {
-            Ok(AdvanceResult::Exhausted)
-        } else {
-            match self.state_mut().buffer.advance() {
-                true => Ok(AdvanceResult::Advanced),
-                false => Ok(AdvanceResult::Waiting),
+        match self.state_mut().buffer.advance() {
+            true => Ok(AdvanceResult::Advanced),
+            false => {
+                if self.is_exhausted() {
+                    Ok(AdvanceResult::Exhausted)
+                } else {
+                    Ok(AdvanceResult::Waiting)
+                }
             }
         }
     }

diff --git a/src/test/cursor.rs b/src/test/cursor.rs
@@ -257,3 +257,34 @@ async fn session_cursor_with_type() {
 
     let _ = cursor_with_type.next(&mut session).await.unwrap().unwrap();
 }
+
+#[cfg_attr(feature = "tokio-runtime", tokio::test)]
+#[cfg_attr(feature = "async-std-runtime", async_std::test)]
+async fn cursor_final_batch() {
+    let client = TestClient::new().await;
+    let coll = client
+        .create_fresh_collection("test_cursor_final_batch", "test", None)
+        .await;
+    coll.insert_many(
+        vec![
+            doc! { "foo": 1 },
+            doc! { "foo": 2 },
+            doc! { "foo": 3 },
+            doc! { "foo": 4 },
+            doc! { "foo": 5 },
+        ],
+        None,
+    )
+    .await
+    .unwrap();
+
+    let mut cursor = coll
+        .find(None, FindOptions::builder().batch_size(3).build())
+        .await
+        .unwrap();
+    let mut found = 0;
+    while cursor.advance().await.unwrap() {
+        found += 1;
+    }
+    assert_eq!(found, 5);
+}