1 files changed, 12 insertions, 18 deletions

diff --git a/src/mongo/s/balancer/migration_manager_test.cpp b/src/mongo/s/balancer/migration_manager_test.cpp
index 645de866c96..bd11a48521f 100644
--- a/src/mongo/s/balancer/migration_manager_test.cpp
+++ b/src/mongo/s/balancer/migration_manager_test.cpp
@@ -155,9 +155,6 @@ protected:
 
     const KeyPattern kKeyPattern = KeyPattern(BSON(kPattern << 1));
 
-    // Cluster identity to pass to the migration manager
-    const OID _clusterIdentity{OID::gen()};
-
     std::unique_ptr<MigrationManager> _migrationManager;
 
 private:
@@ -168,9 +165,8 @@ private:
 void MigrationManagerTest::setUp() {
     ConfigServerTestFixture::setUp();
     _migrationManager = stdx::make_unique<MigrationManager>(getServiceContext());
-    _migrationManager->startRecovery();
-    _migrationManager->finishRecovery(
-        operationContext(), _clusterIdentity, 0, kDefaultSecondaryThrottle, false);
+    _migrationManager->startRecoveryAndAcquireDistLocks(operationContext());
+    _migrationManager->finishRecovery(operationContext(), 0, kDefaultSecondaryThrottle, false);
 }
 
 void MigrationManagerTest::tearDown() {
@@ -834,9 +830,8 @@ TEST_F(MigrationManagerTest, RestartMigrationManager) {
     // Go through the lifecycle of the migration manager
     _migrationManager->interruptAndDisableMigrations();
     _migrationManager->drainActiveMigrations();
-    _migrationManager->startRecovery();
-    _migrationManager->finishRecovery(
-        operationContext(), _clusterIdentity, 0, kDefaultSecondaryThrottle, false);
+    _migrationManager->startRecoveryAndAcquireDistLocks(operationContext());
+    _migrationManager->finishRecovery(operationContext(), 0, kDefaultSecondaryThrottle, false);
 
     auto future = launchAsync([&] {
         Client::initThreadIfNotAlready("Test");
@@ -881,7 +876,6 @@ TEST_F(MigrationManagerTest, MigrationRecovery) {
 
     _migrationManager->interruptAndDisableMigrations();
     _migrationManager->drainActiveMigrations();
-    _migrationManager->startRecovery();
 
     // Set up two fake active migrations by writing documents to the config.migrations collection.
     setUpMigration(collName,
@@ -895,6 +889,8 @@ TEST_F(MigrationManagerTest, MigrationRecovery) {
                    kShardId3.toString(),
                    chunk2.getShard().toString());
 
+    _migrationManager->startRecoveryAndAcquireDistLocks(operationContext());
+
     auto future = launchAsync([this] {
         Client::initThreadIfNotAlready("Test");
         auto txn = cc().makeOperationContext();
@@ -904,8 +900,7 @@ TEST_F(MigrationManagerTest, MigrationRecovery) {
         shardTargeterMock(txn.get(), kShardId0)->setFindHostReturnValue(kShardHost0);
         shardTargeterMock(txn.get(), kShardId2)->setFindHostReturnValue(kShardHost2);
 
-        _migrationManager->finishRecovery(
-            txn.get(), _clusterIdentity, 0, kDefaultSecondaryThrottle, false);
+        _migrationManager->finishRecovery(txn.get(), 0, kDefaultSecondaryThrottle, false);
     });
 
     // Expect two moveChunk commands.
@@ -938,6 +933,9 @@ TEST_F(MigrationManagerTest, FailMigrationRecovery) {
     ChunkType chunk2 =
         setUpChunk(collName, BSON(kPattern << 49), kKeyPattern.globalMax(), kShardId2, version);
 
+    _migrationManager->interruptAndDisableMigrations();
+    _migrationManager->drainActiveMigrations();
+
     // Set up a parsable fake active migration document in the config.migrations collection.
     setUpMigration(collName,
                    chunk1.getMin(),
@@ -945,10 +943,6 @@ TEST_F(MigrationManagerTest, FailMigrationRecovery) {
                    kShardId1.toString(),
                    chunk1.getShard().toString());
 
-    _migrationManager->interruptAndDisableMigrations();
-    _migrationManager->drainActiveMigrations();
-    _migrationManager->startRecovery();
-
     // Set up a fake active migration document that will fail MigrationType parsing -- missing
     // field.
     BSONObjBuilder builder;
@@ -971,8 +965,8 @@ TEST_F(MigrationManagerTest, FailMigrationRecovery) {
         OID::gen(),
         DistLockManager::kSingleLockAttemptTimeout));
 
-    _migrationManager->finishRecovery(
-        operationContext(), _clusterIdentity, 0, kDefaultSecondaryThrottle, false);
+    _migrationManager->startRecoveryAndAcquireDistLocks(operationContext());
+    _migrationManager->finishRecovery(operationContext(), 0, kDefaultSecondaryThrottle, false);
 
     // MigrationManagerTest::tearDown checks that the config.migrations collection is empty and all
     // distributed locks are unlocked.