1 files changed, 21 insertions, 15 deletions

diff --git a/src/mongo/rpc/op_msg_integration_test.cpp b/src/mongo/rpc/op_msg_integration_test.cpp
index 2009a609876..2e5cdad314b 100644
--- a/src/mongo/rpc/op_msg_integration_test.cpp
+++ b/src/mongo/rpc/op_msg_integration_test.cpp
@@ -526,16 +526,18 @@ TEST(OpMsg, ServerHandlesExhaustIsMasterCorrectly) {
     auto fixtureConn = std::unique_ptr<DBClientBase>(
         unittest::getFixtureConnectionString().connect("integration_test", errMsg));
     uassert(ErrorCodes::SocketException, errMsg, fixtureConn);
+    DBClientBase* conn = fixtureConn.get();
 
     // TODO SERVER-44813: Run this test on standalone.
-    // TODO SERVER-44521: Run this test on mongos.
-    if (!fixtureConn->isReplicaSetMember()) {
+    if (!fixtureConn->isReplicaSetMember() && !fixtureConn->isMongos()) {
         return;
     }
 
-    // Connect directly to the primary.
-    DBClientBase* conn = &static_cast<DBClientReplicaSet*>(fixtureConn.get())->masterConn();
-    ASSERT(conn);
+    if (fixtureConn->isReplicaSetMember()) {
+        // Connect directly to the primary.
+        conn = &static_cast<DBClientReplicaSet*>(fixtureConn.get())->masterConn();
+        ASSERT(conn);
+    }
 
     auto tickSource = getGlobalServiceContext()->getTickSource();
 
@@ -592,16 +594,18 @@ TEST(OpMsg, ServerHandlesExhaustIsMasterWithTopologyChange) {
     auto fixtureConn = std::unique_ptr<DBClientBase>(
         unittest::getFixtureConnectionString().connect("integration_test", errMsg));
     uassert(ErrorCodes::SocketException, errMsg, fixtureConn);
+    DBClientBase* conn = fixtureConn.get();
 
     // TODO SERVER-44813: Run this test on standalone.
-    // TODO SERVER-44521: Run this test on mongos.
-    if (!fixtureConn->isReplicaSetMember()) {
+    if (!fixtureConn->isReplicaSetMember() && !fixtureConn->isMongos()) {
         return;
     }
 
-    // Connect directly to the primary.
-    DBClientBase* conn = &static_cast<DBClientReplicaSet*>(fixtureConn.get())->masterConn();
-    ASSERT(conn);
+    if (fixtureConn->isReplicaSetMember()) {
+        // Connect directly to the primary.
+        conn = &static_cast<DBClientReplicaSet*>(fixtureConn.get())->masterConn();
+        ASSERT(conn);
+    }
 
     auto tickSource = getGlobalServiceContext()->getTickSource();
 
@@ -661,16 +665,18 @@ TEST(OpMsg, ServerRejectsExhaustIsMasterWithoutMaxAwaitTimeMS) {
     auto fixtureConn = std::unique_ptr<DBClientBase>(
         unittest::getFixtureConnectionString().connect("integration_test", errMsg));
     uassert(ErrorCodes::SocketException, errMsg, fixtureConn);
+    DBClientBase* conn = fixtureConn.get();
 
     // TODO SERVER-44813: Run this test on standalone.
-    // TODO SERVER-44521: Run this test on mongos.
-    if (!fixtureConn->isReplicaSetMember()) {
+    if (!fixtureConn->isReplicaSetMember() && !fixtureConn->isMongos()) {
         return;
     }
 
-    // Connect directly to the primary.
-    DBClientBase* conn = &static_cast<DBClientReplicaSet*>(fixtureConn.get())->masterConn();
-    ASSERT(conn);
+    if (fixtureConn->isReplicaSetMember()) {
+        // Connect directly to the primary.
+        conn = &static_cast<DBClientReplicaSet*>(fixtureConn.get())->masterConn();
+        ASSERT(conn);
+    }
 
     // Issue an isMaster command with exhaust but no maxAwaitTimeMS.
     auto isMasterCmd = BSON("isMaster" << 1);