SERVER-53444 Make tests that run removeShard in assert.soon to wait for the state to become completed not error on ShardNotFound

1 files changed, 9 insertions, 0 deletions

diff --git a/jstests/sharding/balancing_sessions_collection.js b/jstests/sharding/balancing_sessions_collection.js
index 416eaa6186d..8313b1ca80d 100644
--- a/jstests/sharding/balancing_sessions_collection.js
+++ b/jstests/sharding/balancing_sessions_collection.js
@@ -51,6 +51,15 @@ function removeShardFromCluster(shardName) {
     assert.commandWorked(st.s.adminCommand({removeShard: shardName}));
     assert.soon(function() {
         const res = st.s.adminCommand({removeShard: shardName});
+        if (!res.ok && res.code === ErrorCodes.ShardNotFound) {
+            // If the config server primary steps down right after removing the config.shards doc
+            // for the shard but before responding with "state": "completed", the mongos would retry
+            // the _configsvrRemoveShard command against the new config server primary, which would
+            // not find the removed shard in its ShardRegistry if it has done a ShardRegistry reload
+            // after the config.shards doc for the shard was removed. This would cause the command
+            // to fail with ShardNotFound.
+            return true;
+        }
         assert.commandWorked(res);
         return ("completed" == res.state);
     }, "failed to remove shard " + shardName, kBalancerTimeoutMS);