SERVER-43937 Allow MockDBClientConnection/MockDBServer to mock errors.

4 files changed, 31 insertions, 23 deletions

diff --git a/src/mongo/dbtests/mock/mock_dbclient_connection.cpp b/src/mongo/dbtests/mock/mock_dbclient_connection.cpp
index 59abe29a460..6ad13d3ab3d 100644
--- a/src/mongo/dbtests/mock/mock_dbclient_connection.cpp
+++ b/src/mongo/dbtests/mock/mock_dbclient_connection.cpp
@@ -32,6 +32,7 @@
 #include "mongo/dbtests/mock/mock_dbclient_connection.h"
 
 #include "mongo/client/dbclient_mockcursor.h"
+#include "mongo/rpc/get_status_from_command_result.h"
 #include "mongo/util/net/socket_exception.h"
 #include "mongo/util/time_support.h"
 
@@ -68,12 +69,20 @@ std::pair<rpc::UniqueReply, DBClientBase*> MockDBClientConnection::runCommandWit
     checkConnection();
 
     try {
-        return {_remoteServer->runCommand(_remoteServerInstanceID, request), this};
+        auto reply = _remoteServer->runCommand(_remoteServerInstanceID, request);
+        auto status = getStatusFromCommandResult(reply->getCommandReply());
+        // The real DBClientBase always throws HostUnreachable on network error, so we do the
+        // same here.
+        uassert(ErrorCodes::HostUnreachable,
+                str::stream() << "network error while attempting to run "
+                              << "command '" << request.getCommandName() << "' " << status,
+                !ErrorCodes::isNetworkError(status));
+        return {std::move(reply), this};
     } catch (const mongo::DBException&) {
         _isFailed = true;
         throw;
     }
-}
+}  // namespace mongo
 
 
 std::unique_ptr<mongo::DBClientCursor> MockDBClientConnection::query(
diff --git a/src/mongo/dbtests/mock/mock_remote_db_server.cpp b/src/mongo/dbtests/mock/mock_remote_db_server.cpp
index 6488e5023ff..0f8c0f69d4d 100644
--- a/src/mongo/dbtests/mock/mock_remote_db_server.cpp
+++ b/src/mongo/dbtests/mock/mock_remote_db_server.cpp
@@ -47,20 +47,19 @@ using std::vector;
 
 namespace mongo {
 
-MockRemoteDBServer::CircularBSONIterator::CircularBSONIterator(const vector<BSONObj>& replyVector) {
-    for (std::vector<mongo::BSONObj>::const_iterator iter = replyVector.begin();
-         iter != replyVector.end();
-         ++iter) {
-        _replyObjs.push_back(iter->copy());
+MockRemoteDBServer::CircularBSONIterator::CircularBSONIterator(
+    const vector<StatusWith<BSONObj>>& replyVector) {
+    for (auto iter = replyVector.begin(); iter != replyVector.end(); ++iter) {
+        _replyObjs.push_back(iter->isOK() ? StatusWith(iter->getValue().copy()) : *iter);
     }
 
     _iter = _replyObjs.begin();
 }
 
-BSONObj MockRemoteDBServer::CircularBSONIterator::next() {
+StatusWith<BSONObj> MockRemoteDBServer::CircularBSONIterator::next() {
     verify(_iter != _replyObjs.end());
 
-    BSONObj reply = _iter->copy();
+    StatusWith<BSONObj> reply = _iter->isOK() ? StatusWith(_iter->getValue().copy()) : *_iter;
     ++_iter;
 
     if (_iter == _replyObjs.end()) {
@@ -109,14 +108,15 @@ bool MockRemoteDBServer::isRunning() const {
     return _isRunning;
 }
 
-void MockRemoteDBServer::setCommandReply(const string& cmdName, const mongo::BSONObj& replyObj) {
-    vector<BSONObj> replySequence;
+void MockRemoteDBServer::setCommandReply(const string& cmdName,
+                                         const StatusWith<mongo::BSONObj>& replyObj) {
+    vector<StatusWith<BSONObj>> replySequence;
     replySequence.push_back(replyObj);
     setCommandReply(cmdName, replySequence);
 }
 
 void MockRemoteDBServer::setCommandReply(const string& cmdName,
-                                         const vector<BSONObj>& replySequence) {
+                                         const vector<StatusWith<BSONObj>>& replySequence) {
     scoped_spinlock sLock(_lock);
     _cmdMap[cmdName].reset(new CircularBSONIterator(replySequence));
 }
@@ -146,16 +146,15 @@ rpc::UniqueReply MockRemoteDBServer::runCommand(InstanceID id, const OpMsgReques
     checkIfUp(id);
     std::string cmdName = request.getCommandName().toString();
 
-    BSONObj reply;
-    {
+    StatusWith<BSONObj> reply([this, &cmdName] {
         scoped_spinlock lk(_lock);
 
         uassert(ErrorCodes::IllegalOperation,
                 str::stream() << "no reply for command: " << cmdName,
                 _cmdMap.count(cmdName));
 
-        reply = _cmdMap[cmdName]->next();
-    }
+        return _cmdMap[cmdName]->next();
+    }());
 
     if (_delayMilliSec > 0) {
         mongo::sleepmillis(_delayMilliSec);
diff --git a/src/mongo/dbtests/mock/mock_remote_db_server.h b/src/mongo/dbtests/mock/mock_remote_db_server.h
index 6c0ccc8170d..db6361ac915 100644
--- a/src/mongo/dbtests/mock/mock_remote_db_server.h
+++ b/src/mongo/dbtests/mock/mock_remote_db_server.h
@@ -116,7 +116,7 @@ public:
      * @param cmdName the name of the command
      * @param replyObj the exact reply for the command
      */
-    void setCommandReply(const std::string& cmdName, const mongo::BSONObj& replyObj);
+    void setCommandReply(const std::string& cmdName, const StatusWith<mongo::BSONObj>& replyObj);
 
     /**
      * Sets the reply for a command.
@@ -128,7 +128,7 @@ public:
      *     that requires different results when calling a method.
      */
     void setCommandReply(const std::string& cmdName,
-                         const std::vector<mongo::BSONObj>& replySequence);
+                         const std::vector<StatusWith<mongo::BSONObj>>& replySequence);
 
     /**
      * Inserts a single document to this server.
@@ -203,12 +203,12 @@ private:
         /**
          * Creates a new iterator with a deep copy of the vector.
          */
-        CircularBSONIterator(const std::vector<mongo::BSONObj>& replyVector);
-        mongo::BSONObj next();
+        CircularBSONIterator(const std::vector<StatusWith<mongo::BSONObj>>& replyVector);
+        StatusWith<mongo::BSONObj> next();
 
     private:
-        std::vector<mongo::BSONObj>::iterator _iter;
-        std::vector<mongo::BSONObj> _replyObjs;
+        std::vector<StatusWith<mongo::BSONObj>>::iterator _iter;
+        std::vector<StatusWith<mongo::BSONObj>> _replyObjs;
     };
 
     /**
diff --git a/src/mongo/dbtests/mock_dbclient_conn_test.cpp b/src/mongo/dbtests/mock_dbclient_conn_test.cpp
index 236b80d45d4..936dd270888 100644
--- a/src/mongo/dbtests/mock_dbclient_conn_test.cpp
+++ b/src/mongo/dbtests/mock_dbclient_conn_test.cpp
@@ -411,7 +411,7 @@ TEST(MockDBClientConnTest, CyclingCmd) {
     MockRemoteDBServer server("test");
 
     {
-        vector<BSONObj> isMasterSequence;
+        vector<mongo::StatusWith<BSONObj>> isMasterSequence;
         isMasterSequence.push_back(BSON("set"
                                         << "a"
                                         << "isMaster" << true << "ok" << 1));