address copilot feedback

Author: Almas Abdrazak

driver-reactive-streams/src/main/com/mongodb/reactivestreams/client/internal/ClientSessionPublisherImpl.java

@@ -27,6 +27,7 @@
 import com.mongodb.internal.observability.micrometer.TracingManager;
 import com.mongodb.internal.observability.micrometer.TransactionSpan;
 import com.mongodb.internal.operation.AbortTransactionOperation;
+import com.mongodb.internal.operation.AsyncWriteThenReadOperationCursor;
 import com.mongodb.internal.operation.CommitTransactionOperation;
 import com.mongodb.internal.operation.ReadOperation;
 import com.mongodb.internal.operation.WriteConcernHelper;
@@ -88,7 +89,7 @@ public boolean notifyMessageSent() {
 
     @Override
     public void notifyOperationInitiated(final Object operation) {
-        assertTrue(operation instanceof ReadOperation || operation instanceof WriteOperation);
+        assertTrue(operation instanceof ReadOperation || operation instanceof WriteOperation || operation instanceof AsyncWriteThenReadOperationCursor);
         if (!(hasActiveTransaction() || operation instanceof CommitTransactionOperation)) {
             assertTrue(getPinnedServerAddress() == null
                     || (transactionState != TransactionState.ABORTED && transactionState != TransactionState.NONE));

driver-reactive-streams/src/main/com/mongodb/reactivestreams/client/internal/OperationExecutorImpl.java

@@ -188,6 +188,10 @@ public <T> Mono<AsyncBatchCursor<T>> execute(final AsyncWriteThenReadOperationCu
         notNull("operation", operation);
         notNull("readConcern", readConcern);
 
+        if (session != null) {
+            session.notifyOperationInitiated(operation);
+        }
+
         return Mono.from(subscriber ->
                 clientSessionHelper.withClientSession(session, this)
                         .flatMap(actualClientSession -> {

driver-reactive-streams/src/test/unit/com/mongodb/reactivestreams/client/internal/MapReducePublisherImplTest.java

@@ -20,10 +20,10 @@
 import com.mongodb.ReadPreference;
 import com.mongodb.WriteConcern;
 import com.mongodb.client.model.Sorts;
+import com.mongodb.internal.operation.AsyncWriteThenReadOperationCursor;
 import com.mongodb.internal.operation.MapReduceStatistics;
 import com.mongodb.internal.operation.MapReduceToCollectionOperation;
 import com.mongodb.internal.operation.MapReduceWithInlineResultsOperation;
-
 import com.mongodb.reactivestreams.client.MapReducePublisher;
 
 
@@ -53,6 +53,27 @@ public class MapReducePublisherImplTest extends TestHelper {
     private static final String REDUCE_FUNCTION = "reduceFunction(){}";
     private static final String FINALIZE_FUNCTION = "finalizeFunction(){}";
 
+    @DisplayName("Non-inline MapReduce routes through the write-then-read executor path")
+    @Test
+    void shouldRouteNonInlineMapReduceThroughWriteThenReadPath() {
+        configureBatchCursor();
+        TestOperationExecutor executor = createOperationExecutor(asList(getBatchCursor()));
+
+        MapReducePublisher<Document> publisher =
+                new MapReducePublisherImpl<>(null, createMongoOperationPublisher(executor),
+                        MAP_FUNCTION, REDUCE_FUNCTION)
+                        .collectionName(NAMESPACE.getCollectionName());
+
+        Flux.from(publisher).blockFirst();
+
+        AsyncWriteThenReadOperationCursor<?> op = executor.getWriteThenReadOperation();
+        assertNotNull(op, "expected a write-then-read operation");
+        assertEquals(NAMESPACE, op.getNamespace());
+        // Must not fall through to the plain read-operation path.
+        assertNull(executor.getReadOperation());
+    }
+
+
     @DisplayName("Inline MapReduce still routes through the read-operation path")
     @Test
     void shouldRouteInlineMapReduceThroughReadOperationPath() {