Add CommandMessageTest.getCommandDocumentFromClientBulkWrite

JAVA-5529

Author: stIncMale

driver-core/src/main/com/mongodb/internal/operation/ClientBulkWriteOperation.java

@@ -39,6 +39,7 @@
 import com.mongodb.client.model.bulk.ClientUpdateResult;
 import com.mongodb.connection.ConnectionDescription;
 import com.mongodb.internal.TimeoutContext;
+import com.mongodb.internal.VisibleForTesting;
 import com.mongodb.internal.async.function.RetryState;
 import com.mongodb.internal.binding.ConnectionSource;
 import com.mongodb.internal.binding.WriteBinding;
@@ -109,6 +110,8 @@
 import static com.mongodb.assertions.Assertions.assertNotNull;
 import static com.mongodb.assertions.Assertions.assertTrue;
 import static com.mongodb.assertions.Assertions.fail;
+import static com.mongodb.internal.VisibleForTesting.AccessModifier.PACKAGE;
+import static com.mongodb.internal.VisibleForTesting.AccessModifier.PRIVATE;
 import static com.mongodb.internal.connection.BsonWriterHelper.decorateWithDocumentSizeChecking;
 import static com.mongodb.internal.connection.DualMessageSequences.WritersProviderAndLimitsChecker.WriteResult.FAIL_LIMIT_EXCEEDED;
 import static com.mongodb.internal.connection.DualMessageSequences.WritersProviderAndLimitsChecker.WriteResult.OK_LIMIT_NOT_REACHED;
@@ -682,7 +685,8 @@ public static final class OpsAndNsInfo extends DualMessageSequences {
             private final ConcreteClientBulkWriteOptions options;
             private final Supplier<Long> doIfCommandIsRetryableAndAdvanceGetTxnNumber;
 
-            OpsAndNsInfo(
+            @VisibleForTesting(otherwise = PACKAGE)
+            public OpsAndNsInfo(
                     final boolean effectiveRetryWrites,
                     final List<? extends ClientNamespacedWriteModel> models,
                     @Nullable final Integer maxStoredDocumentSize,
@@ -941,10 +945,12 @@ Map<Integer, BsonValue> getInsertModelDocumentIds() {
     /**
      * Exactly one instance must be used per {@linkplain #executeBatch(int, WriteConcern, WriteBinding, ResultAccumulator) batch}.
      */
-    private final class BatchEncoder {
+    @VisibleForTesting(otherwise = PRIVATE)
+    public final class BatchEncoder {
         private EncodedBatchInfo encodedBatchInfo;
 
-        BatchEncoder() {
+        @VisibleForTesting(otherwise = PACKAGE)
+        public BatchEncoder() {
             encodedBatchInfo = new EncodedBatchInfo();
         }
 

driver-core/src/test/unit/com/mongodb/internal/connection/CommandMessageSpecification.groovy

@@ -46,6 +46,9 @@ import static com.mongodb.internal.connection.SplittablePayload.Type.INSERT
 import static com.mongodb.internal.operation.ServerVersionHelper.FOUR_DOT_ZERO_WIRE_VERSION
 import static com.mongodb.internal.operation.ServerVersionHelper.LATEST_WIRE_VERSION
 
+/**
+ * New tests must be added to {@link CommandMessageTest}.
+ */
 class CommandMessageSpecification extends Specification {
 
     def namespace = new MongoNamespace('db.test')

driver-core/src/test/unit/com/mongodb/internal/connection/CommandMessageTest.java

@@ -20,19 +20,39 @@
 import com.mongodb.MongoOperationTimeoutException;
 import com.mongodb.ReadConcern;
 import com.mongodb.ReadPreference;
+import com.mongodb.WriteConcern;
+import com.mongodb.client.model.bulk.ClientNamespacedWriteModel;
 import com.mongodb.connection.ClusterConnectionMode;
 import com.mongodb.connection.ServerType;
+import com.mongodb.internal.IgnorableRequestContext;
 import com.mongodb.internal.TimeoutContext;
+import com.mongodb.internal.TimeoutSettings;
+import com.mongodb.internal.client.model.bulk.ConcreteClientBulkWriteOptions;
 import com.mongodb.internal.connection.MessageSequences.EmptyMessageSequences;
+import com.mongodb.internal.operation.ClientBulkWriteOperation;
+import com.mongodb.internal.operation.ClientBulkWriteOperation.ClientBulkWriteCommand.OpsAndNsInfo;
 import com.mongodb.internal.session.SessionContext;
 import com.mongodb.internal.validator.NoOpFieldNameValidator;
+import org.bson.BsonArray;
+import org.bson.BsonBoolean;
 import org.bson.BsonDocument;
+import org.bson.BsonInt32;
 import org.bson.BsonString;
 import org.bson.BsonTimestamp;
 import org.junit.jupiter.api.Test;
 
+import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static com.mongodb.MongoClientSettings.getDefaultCodecRegistry;
+import static com.mongodb.client.model.bulk.ClientBulkWriteOptions.clientBulkWriteOptions;
 import static com.mongodb.internal.mockito.MongoMockito.mock;
 import static com.mongodb.internal.operation.ServerVersionHelper.FOUR_DOT_ZERO_WIRE_VERSION;
+import static com.mongodb.internal.operation.ServerVersionHelper.LATEST_WIRE_VERSION;
+import static java.util.Arrays.asList;
+import static java.util.Collections.singletonList;
+import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doThrow;
@@ -105,4 +125,49 @@ void encodeShouldNotAddExtraElementsFromTimeoutContextWhenConnectedToMongoCrypt(
             verifyNoInteractions(timeoutContext);
         }
     }
+
+    @Test
+    void getCommandDocumentFromClientBulkWrite() {
+        MongoNamespace ns = new MongoNamespace("db", "test");
+        boolean retryWrites = false;
+        BsonDocument command = new BsonDocument("bulkWrite", new BsonInt32(1))
+                .append("errorsOnly", BsonBoolean.valueOf(false))
+                .append("ordered", BsonBoolean.valueOf(true));
+        List<BsonDocument> documents = IntStream.range(0, 2).mapToObj(i -> new BsonDocument("_id", new BsonInt32(i)))
+                .collect(Collectors.toList());
+        List<ClientNamespacedWriteModel> writeModels = asList(
+                ClientNamespacedWriteModel.insertOne(ns, documents.get(0)),
+                ClientNamespacedWriteModel.insertOne(ns, documents.get(1)));
+        OpsAndNsInfo opsAndNsInfo = new OpsAndNsInfo(
+                retryWrites,
+                writeModels,
+                null,
+                new ClientBulkWriteOperation(
+                        writeModels,
+                        clientBulkWriteOptions(),
+                        WriteConcern.MAJORITY,
+                        retryWrites,
+                        getDefaultCodecRegistry()
+                ).new BatchEncoder(),
+                (ConcreteClientBulkWriteOptions) clientBulkWriteOptions(),
+                () -> 1L);
+        BsonDocument expectedCommandDocument = command.clone()
+                .append("$db", new BsonString(ns.getDatabaseName()))
+                .append("ops", new BsonArray(asList(
+                        new BsonDocument("insert", new BsonInt32(0)).append("document", documents.get(0)),
+                        new BsonDocument("insert", new BsonInt32(0)).append("document", documents.get(1)))))
+                .append("nsInfo", new BsonArray(singletonList(new BsonDocument("ns", new BsonString(ns.toString())))));
+        CommandMessage commandMessage = new CommandMessage(
+                ns, command, NoOpFieldNameValidator.INSTANCE, ReadPreference.primary(),
+                MessageSettings.builder().maxWireVersion(LATEST_WIRE_VERSION).build(), true, opsAndNsInfo, ClusterConnectionMode.MULTIPLE, null);
+        try (ByteBufferBsonOutput output = new ByteBufferBsonOutput(new SimpleBufferProvider())) {
+            commandMessage.encode(
+                    output,
+                    new OperationContext(
+                            IgnorableRequestContext.INSTANCE, NoOpSessionContext.INSTANCE,
+                            new TimeoutContext(TimeoutSettings.DEFAULT), null));
+            BsonDocument actualCommandDocument = commandMessage.getCommandDocument(output);
+            assertEquals(expectedCommandDocument, actualCommandDocument);
+        }
+    }
 }