SERVER-29876 ReplicationCoordinatorMock::AwaitReplicationReturnValueFunction accepts optime

3 files changed, 9 insertions, 7 deletions

diff --git a/src/mongo/db/catalog/drop_database_test.cpp b/src/mongo/db/catalog/drop_database_test.cpp
index bbb07aaa1ed..ee552ad1c78 100644
--- a/src/mongo/db/catalog/drop_database_test.cpp
+++ b/src/mongo/db/catalog/drop_database_test.cpp
@@ -308,7 +308,7 @@ void _testDropDatabaseResetsDropPendingStateIfAwaitReplicationFails(OperationCon
 TEST_F(DropDatabaseTest,
        DropDatabaseResetsDropPendingStateIfAwaitReplicationFailsAndDatabaseIsPresent) {
     // Update ReplicationCoordinatorMock so that awaitReplicationOfLastOpForClient() fails.
-    _replCoord->setAwaitReplicationReturnValueFunction([] {
+    _replCoord->setAwaitReplicationReturnValueFunction([](const repl::OpTime&) {
         return repl::ReplicationCoordinator::StatusAndDuration(
             Status(ErrorCodes::WriteConcernFailed, ""), Milliseconds(0));
     });
@@ -319,7 +319,7 @@ TEST_F(DropDatabaseTest,
 TEST_F(DropDatabaseTest,
        DropDatabaseResetsDropPendingStateIfAwaitReplicationFailsAndDatabaseIsMissing) {
     // Update ReplicationCoordinatorMock so that awaitReplicationOfLastOpForClient() fails.
-    _replCoord->setAwaitReplicationReturnValueFunction([this] {
+    _replCoord->setAwaitReplicationReturnValueFunction([this](const repl::OpTime&) {
         _removeDatabaseFromCatalog(_opCtx.get(), _nss.db());
         return repl::ReplicationCoordinator::StatusAndDuration(
             Status(ErrorCodes::WriteConcernFailed, ""), Milliseconds(0));
@@ -331,7 +331,7 @@ TEST_F(DropDatabaseTest,
 TEST_F(DropDatabaseTest,
        DropDatabaseReturnsNamespaceNotFoundIfDatabaseIsRemovedAfterCollectionsDropsAreReplicated) {
     // Update ReplicationCoordinatorMock so that awaitReplicationOfLastOpForClient() fails.
-    _replCoord->setAwaitReplicationReturnValueFunction([this] {
+    _replCoord->setAwaitReplicationReturnValueFunction([this](const repl::OpTime&) {
         _removeDatabaseFromCatalog(_opCtx.get(), _nss.db());
         return repl::ReplicationCoordinator::StatusAndDuration(Status::OK(), Milliseconds(0));
     });
diff --git a/src/mongo/db/repl/replication_coordinator_mock.cpp b/src/mongo/db/repl/replication_coordinator_mock.cpp
index db032d0b64a..02097d19815 100644
--- a/src/mongo/db/repl/replication_coordinator_mock.cpp
+++ b/src/mongo/db/repl/replication_coordinator_mock.cpp
@@ -113,13 +113,13 @@ void ReplicationCoordinatorMock::clearSyncSourceBlacklist() {}
 
 ReplicationCoordinator::StatusAndDuration ReplicationCoordinatorMock::awaitReplication(
     OperationContext* opCtx, const OpTime& opTime, const WriteConcernOptions& writeConcern) {
-    return _awaitReplicationReturnValueFunction();
+    return _awaitReplicationReturnValueFunction(opTime);
 }
 
 ReplicationCoordinator::StatusAndDuration
 ReplicationCoordinatorMock::awaitReplicationOfLastOpForClient(
     OperationContext* opCtx, const WriteConcernOptions& writeConcern) {
-    return _awaitReplicationReturnValueFunction();
+    return _awaitReplicationReturnValueFunction({});
 }
 
 void ReplicationCoordinatorMock::setAwaitReplicationReturnValueFunction(
diff --git a/src/mongo/db/repl/replication_coordinator_mock.h b/src/mongo/db/repl/replication_coordinator_mock.h
index 64ec28f6d9e..3ae99bb0dc4 100644
--- a/src/mongo/db/repl/replication_coordinator_mock.h
+++ b/src/mongo/db/repl/replication_coordinator_mock.h
@@ -280,8 +280,10 @@ public:
     /**
      * Sets the function to generate the return value for calls to awaitReplication() and
      * awaitReplicationOfLastOpForClient().
+     * 'opTime' is the optime passed to awaitReplication() and set to null when called from
+     * awaitReplicationOfLastOpForClient().
      */
-    using AwaitReplicationReturnValueFunction = stdx::function<StatusAndDuration()>;
+    using AwaitReplicationReturnValueFunction = stdx::function<StatusAndDuration(const OpTime&)>;
     void setAwaitReplicationReturnValueFunction(
         AwaitReplicationReturnValueFunction returnValueFunction);
 
@@ -306,7 +308,7 @@ private:
     OpTime _myLastDurableOpTime;
     OpTime _myLastAppliedOpTime;
     ReplSetConfig _getConfigReturnValue;
-    AwaitReplicationReturnValueFunction _awaitReplicationReturnValueFunction = [] {
+    AwaitReplicationReturnValueFunction _awaitReplicationReturnValueFunction = [](const OpTime&) {
         return StatusAndDuration(Status::OK(), Milliseconds(0));
     };
     bool _alwaysAllowWrites = false;