add batching prose tests

Author: isabelatkinson

Diff for: src/test/bulk_write.rs

@@ -14,57 +14,103 @@ async fn run_unified() {
 }
 
 #[tokio::test]
-async fn command_batching() {
+async fn max_write_batch_size_batching() {
     let handler = Arc::new(EventHandler::new());
     let client = Client::test_builder()
         .event_handler(handler.clone())
         .build()
         .await;
     let mut subscriber = handler.subscribe();
 
-    let max_object_size = client.server_info.max_bson_object_size as usize;
-    let max_message_size = client.server_info.max_message_size_bytes as usize;
+    let max_write_batch_size = client.server_info.max_write_batch_size.unwrap() as usize;
 
-    let namespace = Namespace::new("command_batching", "command_batching");
-    let large_doc = doc! {"a": "b".repeat(max_object_size / 2)};
-    let models = vec![
-        WriteModel::InsertOne {
-            namespace: namespace.clone(),
-            document: large_doc,
-        };
-        3
-    ];
-    client.bulk_write(models).await.unwrap();
+    let model = WriteModel::InsertOne {
+        namespace: Namespace::new("db", "coll"),
+        document: doc! { "a": "b" },
+    };
+    let models = vec![model; max_write_batch_size + 1];
+
+    let result = client.bulk_write(models).await.unwrap();
+    assert_eq!(result.inserted_count as usize, max_write_batch_size + 1);
+
+    let (first_started, _) = subscriber
+        .wait_for_successful_command_execution(Duration::from_millis(500), "bulkWrite")
+        .await
+        .expect("no events observed");
+    let first_len = first_started.command.get_array("ops").unwrap().len();
+    assert_eq!(first_len, max_write_batch_size);
+
+    let (second_started, _) = subscriber
+        .wait_for_successful_command_execution(Duration::from_millis(500), "bulkWrite")
+        .await
+        .expect("no events observed");
+    let second_len = second_started.command.get_array("ops").unwrap().len();
+    assert_eq!(second_len, 1);
+}
+
+#[tokio::test]
+async fn max_bson_object_size_with_document_sequences() {
+    let handler = Arc::new(EventHandler::new());
+    let client = Client::test_builder()
+        .event_handler(handler.clone())
+        .build()
+        .await;
+    let mut subscriber = handler.subscribe();
+
+    let max_bson_object_size = client.server_info.max_bson_object_size as usize;
+
+    let document = doc! { "a": "b".repeat(max_bson_object_size / 2) };
+    let model = WriteModel::InsertOne {
+        namespace: Namespace::new("db", "coll"),
+        document,
+    };
+    let models = vec![model; 2];
+
+    let result = client.bulk_write(models).await.unwrap();
+    assert_eq!(result.inserted_count as usize, 2);
 
     let (started, _) = subscriber
         .wait_for_successful_command_execution(Duration::from_millis(500), "bulkWrite")
         .await
         .expect("no events observed");
-    let ops = started.command.get_array("ops").unwrap();
-    assert_eq!(ops.len(), 3);
-
-    let large_doc = doc! { "a": "b".repeat(max_object_size - 5000) };
-    let num_models = max_message_size / max_object_size + 1;
-    let models = vec![
-        WriteModel::InsertOne {
-            namespace: namespace.clone(),
-            document: large_doc
-        };
-        num_models
-    ];
-    client.bulk_write(models).await.unwrap();
+    let len = started.command.get_array("ops").unwrap().len();
+    assert_eq!(len, 2);
+}
+
+#[tokio::test]
+async fn max_message_size_bytes_batching() {
+    let handler = Arc::new(EventHandler::new());
+    let client = Client::test_builder()
+        .event_handler(handler.clone())
+        .build()
+        .await;
+    let mut subscriber = handler.subscribe();
+
+    let max_bson_object_size = client.server_info.max_bson_object_size as usize;
+    let max_message_size_bytes = client.server_info.max_message_size_bytes as usize;
+
+    let document = doc! { "a": "b".repeat(max_bson_object_size - 500) };
+    let model = WriteModel::InsertOne {
+        namespace: Namespace::new("db", "coll"),
+        document,
+    };
+    let num_models = max_message_size_bytes / max_bson_object_size + 1;
+    let models = vec![model; num_models];
+
+    let result = client.bulk_write(models).await.unwrap();
+    assert_eq!(result.inserted_count as usize, num_models);
 
     let (first_started, _) = subscriber
         .wait_for_successful_command_execution(Duration::from_millis(500), "bulkWrite")
         .await
         .expect("no events observed");
-    let first_len = first_started.command.get_array("ops").unwrap().len();
-    assert!(first_len < num_models);
+    let first_ops_len = first_started.command.get_array("ops").unwrap().len();
 
     let (second_started, _) = subscriber
         .wait_for_successful_command_execution(Duration::from_millis(500), "bulkWrite")
         .await
         .expect("no events observed");
-    let second_len = second_started.command.get_array("ops").unwrap().len();
-    assert_eq!(first_len + second_len, num_models);
+    let second_ops_len = second_started.command.get_array("ops").unwrap().len();
+
+    assert_eq!(first_ops_len + second_ops_len, num_models);
 }