1 files changed, 42 insertions, 2 deletions

diff --git a/src/mongo/executor/network_test_env.cpp b/src/mongo/executor/network_test_env.cpp
index 64241e72c12..21ad1ef1c6b 100644
--- a/src/mongo/executor/network_test_env.cpp
+++ b/src/mongo/executor/network_test_env.cpp
@@ -69,9 +69,22 @@ void NetworkTestEnv::onCommandWithMetadata(OnCommandWithMetadataFunction func) {
 
     const NetworkInterfaceMock::NetworkOperationIterator noi = _mockNetwork->getNextReadyRequest();
     const RemoteCommandRequest& request = noi->getRequest();
-    _mockNetwork->scheduleResponse(noi, _mockNetwork->now(), func(request));
-    _mockNetwork->runReadyNetworkOperations();
 
+    const auto cmdResponseStatus = func(request);
+    const auto cmdResponse = cmdResponseStatus.getValue();
+
+    BSONObjBuilder result;
+
+    if (cmdResponseStatus.isOK()) {
+        result.appendElements(cmdResponse.data);
+    }
+
+    Command::appendCommandStatus(result, cmdResponseStatus.getStatus());
+
+    const RemoteCommandResponse response(result.obj(), cmdResponse.metadata, Milliseconds(1));
+
+    _mockNetwork->scheduleResponse(noi, _mockNetwork->now(), response);
+    _mockNetwork->runReadyNetworkOperations();
     _mockNetwork->exitNetwork();
 }
 
@@ -97,5 +110,32 @@ void NetworkTestEnv::onFindCommand(OnFindCommandFunction func) {
     });
 }
 
+void NetworkTestEnv::onFindWithMetadataCommand(OnFindCommandWithMetadataFunction func) {
+    onCommandWithMetadata(
+        [&func](const RemoteCommandRequest& request) -> StatusWith<RemoteCommandResponse> {
+            const auto& resultStatus = func(request);
+
+            if (!resultStatus.isOK()) {
+                return resultStatus.getStatus();
+            }
+
+            std::vector<BSONObj> result;
+            BSONObj metadata;
+            std::tie(result, metadata) = resultStatus.getValue();
+
+            BSONArrayBuilder arr;
+            for (const auto& obj : result) {
+                arr.append(obj);
+            }
+
+            const NamespaceString nss =
+                NamespaceString(request.dbname, request.cmdObj.firstElement().String());
+            BSONObjBuilder resultBuilder;
+            appendCursorResponseObject(0LL, nss.toString(), arr.arr(), &resultBuilder);
+
+            return RemoteCommandResponse(resultBuilder.obj(), metadata, Milliseconds(1));
+        });
+}
+
 }  // namespace executor
 }  // namespace mongo