SERVER-20318 Allow NetworkTestEnv::OnCommand to set non ok status to response

1 files changed, 12 insertions, 10 deletions

diff --git a/src/mongo/executor/network_test_env.cpp b/src/mongo/executor/network_test_env.cpp
index c1639cf7822..fae46697d5e 100644
--- a/src/mongo/executor/network_test_env.cpp
+++ b/src/mongo/executor/network_test_env.cpp
@@ -53,13 +53,14 @@ void NetworkTestEnv::onCommand(OnCommandFunction func) {
 
     if (resultStatus.isOK()) {
         result.appendElements(resultStatus.getValue());
-    }
-
-    Command::appendCommandStatus(result, resultStatus.getStatus());
+        Command::appendCommandStatus(result, resultStatus.getStatus());
+        const RemoteCommandResponse response(result.obj(), BSONObj(), Milliseconds(1));
 
-    const RemoteCommandResponse response(result.obj(), BSONObj(), Milliseconds(1));
+        _mockNetwork->scheduleResponse(noi, _mockNetwork->now(), response);
+    } else {
+        _mockNetwork->scheduleResponse(noi, _mockNetwork->now(), resultStatus.getStatus());
+    }
 
-    _mockNetwork->scheduleResponse(noi, _mockNetwork->now(), response);
     _mockNetwork->runReadyNetworkOperations();
     _mockNetwork->exitNetwork();
 }
@@ -77,13 +78,14 @@ void NetworkTestEnv::onCommandWithMetadata(OnCommandWithMetadataFunction func) {
 
     if (cmdResponseStatus.isOK()) {
         result.appendElements(cmdResponse.data);
-    }
-
-    Command::appendCommandStatus(result, cmdResponseStatus.getStatus());
+        Command::appendCommandStatus(result, cmdResponseStatus.getStatus());
+        const RemoteCommandResponse response(result.obj(), cmdResponse.metadata, Milliseconds(1));
 
-    const RemoteCommandResponse response(result.obj(), cmdResponse.metadata, Milliseconds(1));
+        _mockNetwork->scheduleResponse(noi, _mockNetwork->now(), response);
+    } else {
+        _mockNetwork->scheduleResponse(noi, _mockNetwork->now(), cmdResponseStatus.getStatus());
+    }
 
-    _mockNetwork->scheduleResponse(noi, _mockNetwork->now(), response);
     _mockNetwork->runReadyNetworkOperations();
     _mockNetwork->exitNetwork();
 }