1 files changed, 23 insertions, 26 deletions

diff --git a/jstests/replsets/tenant_migration_donor_retry.js b/jstests/replsets/tenant_migration_donor_retry.js
index 281c7a22412..628ba16b927 100644
--- a/jstests/replsets/tenant_migration_donor_retry.js
+++ b/jstests/replsets/tenant_migration_donor_retry.js
@@ -22,30 +22,31 @@ load("jstests/libs/uuid_util.js");
 load("jstests/replsets/libs/tenant_migration_test.js");
 load("jstests/replsets/libs/tenant_migration_util.js");
 
-const kGarbageCollectionDelayMS = 5 * 1000;
 const kTenantIdPrefix = "testTenantId";
 let testNum = 0;
 
-const garbageCollectionOpts = {
-    // Set the delay before a donor state doc is garbage collected to be short to speed
-    // up the test.
-    tenantMigrationGarbageCollectionDelayMS: kGarbageCollectionDelayMS,
-    ttlMonitorSleepSecs: 1
-};
-
 function setup() {
     const donorRst = new ReplSetTest({
         name: "donorRst",
         nodes: 1,
-        nodeOptions: Object.assign(TenantMigrationUtil.makeX509OptionsForTest().donor,
-                                   {setParameter: garbageCollectionOpts})
+        nodeOptions: Object.assign(TenantMigrationUtil.makeX509OptionsForTest().donor, {
+            setParameter: {
+                // Set the delay before a donor state doc is garbage collected to be short to speed
+                // up the test.
+                tenantMigrationGarbageCollectionDelayMS: 0,
+                ttlMonitorSleepSecs: 1
+            }
+        })
     });
 
     donorRst.startSet();
     donorRst.initiate();
 
-    const tenantMigrationTest = new TenantMigrationTest(
-        {name: jsTestName(), donorRst: donorRst, quickGarbageCollection: true});
+    const tenantMigrationTest = new TenantMigrationTest({
+        name: jsTestName(),
+        donorRst: donorRst,
+        quickGarbageCollection: true,
+    });
     return {
         tenantMigrationTest,
         teardown: function() {
@@ -94,7 +95,6 @@ function testDonorRetryRecipientSyncDataCmdOnError(tenantMigrationTest, errorCod
 
     TenantMigrationTest.assertCommitted(
         tenantMigrationTest.waitForMigrationToComplete(migrationOpts));
-    assert.commandWorked(tenantMigrationTest.forgetMigration(migrationOpts.migrationIdString));
 
     return migrationId;
 }
@@ -122,7 +122,8 @@ function testDonorRetryRecipientForgetMigrationCmdOnError(tenantMigrationTest, e
                                 },
                                 {times: 1});
 
-    TenantMigrationTest.assertCommitted(tenantMigrationTest.runMigration(migrationOpts));
+    TenantMigrationTest.assertCommitted(
+        tenantMigrationTest.runMigration(migrationOpts, {automaticForgetMigration: false}));
 
     // Verify that the initial recipientForgetMigration command failed.
     assert.commandWorked(tenantMigrationTest.forgetMigration(migrationOpts.migrationIdString));
@@ -145,6 +146,7 @@ function testDonorRetryRecipientForgetMigrationCmdOnError(tenantMigrationTest, e
         tenantMigrationTest.getDonorPrimary().getCollection(TenantMigrationTest.kConfigDonorsNS);
     assert.eq(TenantMigrationTest.DonorState.kCommitted,
               configDonorsColl.findOne({_id: migrationId}).state);
+    assert.commandWorked(tenantMigrationTest.forgetMigration(extractUUIDFromObject(migrationId)));
     teardown();
 })();
 
@@ -160,6 +162,7 @@ function testDonorRetryRecipientForgetMigrationCmdOnError(tenantMigrationTest, e
         tenantMigrationTest.getDonorPrimary().getCollection(TenantMigrationTest.kConfigDonorsNS);
     assert.eq(TenantMigrationTest.DonorState.kCommitted,
               configDonorsColl.findOne({_id: migrationId}).state);
+    assert.commandWorked(tenantMigrationTest.forgetMigration(extractUUIDFromObject(migrationId)));
     teardown();
 })();
 
@@ -190,6 +193,7 @@ function testDonorRetryRecipientForgetMigrationCmdOnError(tenantMigrationTest, e
         tenantMigrationTest.getDonorPrimary().getCollection(TenantMigrationTest.kConfigDonorsNS);
     assert.eq(TenantMigrationTest.DonorState.kCommitted,
               configDonorsColl.findOne({_id: migrationId}).state);
+    assert.commandWorked(tenantMigrationTest.forgetMigration(extractUUIDFromObject(migrationId)));
     teardown();
 })();
 
@@ -272,7 +276,7 @@ const kWriteErrorTimeMS = 50;
     jsTest.log("Test that the donor retries state doc update on retriable errors");
 
     const {tenantMigrationTest, teardown} = setup();
-    const tenantId = kTenantIdPrefix + "RetryOnStateDocUpdateError";
+    const tenantId = `${kTenantIdPrefix}RetryOnStateDocUpdateError`;
 
     const migrationId = UUID();
     const migrationOpts = {
@@ -307,17 +311,10 @@ const kWriteErrorTimeMS = 50;
     fp.off();
     migrationThread.join();
 
-    const donorStateDoc = tenantMigrationTest.getDonorPrimary()
-                              .getCollection(TenantMigrationTest.kConfigDonorsNS)
-                              .findOne({_id: migrationId});
-    assert.eq(donorStateDoc.state, TenantMigrationTest.DonorState.kCommitted);
-    assert(donorStateDoc.expireAt);
-
-    // Check that the recipient state doc is also correctly marked as garbage collectable.
-    const recipientStateDoc = tenantMigrationTest.getRecipientPrimary()
-                                  .getCollection(TenantMigrationTest.kConfigRecipientsNS)
-                                  .findOne({_id: migrationId});
-    assert(recipientStateDoc.expireAt);
+    // The state docs will only be completed and marked as garbage collectable if the
+    // update succeeds.
+    tenantMigrationTest.waitForMigrationGarbageCollection(migrationId, tenantId);
+
     teardown();
 })();
 })();