PYTHON-4203 Update prose tests for mongos deprioritization during retryable ops (#2430)

Author: blink1073

test/asynchronous/test_retryable_reads.py

@@ -21,7 +21,7 @@
 import threading
 from test.asynchronous.utils import async_set_fail_point
 
-from pymongo.errors import AutoReconnect
+from pymongo.errors import OperationFailure
 
 sys.path[0:0] = [""]
 
@@ -147,15 +147,11 @@ async def test_pool_paused_error_is_retryable(self):
 class TestRetryableReads(AsyncIntegrationTest):
     @async_client_context.require_multiple_mongoses
     @async_client_context.require_failCommand_fail_point
-    async def test_retryable_reads_in_sharded_cluster_multiple_available(self):
+    async def test_retryable_reads_are_retried_on_a_different_mongos_when_one_is_available(self):
         fail_command = {
             "configureFailPoint": "failCommand",
             "mode": {"times": 1},
-            "data": {
-                "failCommands": ["find"],
-                "closeConnection": True,
-                "appName": "retryableReadTest",
-            },
+            "data": {"failCommands": ["find"], "errorCode": 6},
         }
 
         mongos_clients = []
@@ -168,12 +164,11 @@ async def test_retryable_reads_in_sharded_cluster_multiple_available(self):
         listener = OvertCommandListener()
         client = await self.async_rs_or_single_client(
             async_client_context.mongos_seeds(),
-            appName="retryableReadTest",
             event_listeners=[listener],
             retryReads=True,
         )
 
-        with self.assertRaises(AutoReconnect):
+        with self.assertRaises(OperationFailure):
             await client.t.t.find_one({})
 
         # Disable failpoints on each mongos
@@ -184,6 +179,45 @@ async def test_retryable_reads_in_sharded_cluster_multiple_available(self):
         self.assertEqual(len(listener.failed_events), 2)
         self.assertEqual(len(listener.succeeded_events), 0)
 
+        #  Assert that both events occurred on different mongos.
+        assert listener.failed_events[0].connection_id != listener.failed_events[1].connection_id
+
+    @async_client_context.require_multiple_mongoses
+    @async_client_context.require_failCommand_fail_point
+    async def test_retryable_reads_are_retried_on_the_same_mongos_when_no_others_are_available(
+        self
+    ):
+        fail_command = {
+            "configureFailPoint": "failCommand",
+            "mode": {"times": 1},
+            "data": {"failCommands": ["find"], "errorCode": 6},
+        }
+
+        host = async_client_context.mongos_seeds().split(",")[0]
+        mongos_client = await self.async_rs_or_single_client(host)
+        await async_set_fail_point(mongos_client, fail_command)
+
+        listener = OvertCommandListener()
+        client = await self.async_rs_or_single_client(
+            host,
+            directConnection=False,
+            event_listeners=[listener],
+            retryReads=True,
+        )
+
+        await client.t.t.find_one({})
+
+        # Disable failpoint.
+        fail_command["mode"] = "off"
+        await async_set_fail_point(mongos_client, fail_command)
+
+        # Assert that exactly one failed command event and one succeeded command event occurred.
+        self.assertEqual(len(listener.failed_events), 1)
+        self.assertEqual(len(listener.succeeded_events), 1)
+
+        #  Assert that both events occurred on the same mongos.
+        assert listener.succeeded_events[0].connection_id == listener.failed_events[0].connection_id
+
 
 if __name__ == "__main__":
     unittest.main()

test/test_retryable_reads.py

@@ -21,7 +21,7 @@
 import threading
 from test.utils import set_fail_point
 
-from pymongo.errors import AutoReconnect
+from pymongo.errors import OperationFailure
 
 sys.path[0:0] = [""]
 
@@ -147,15 +147,11 @@ def test_pool_paused_error_is_retryable(self):
 class TestRetryableReads(IntegrationTest):
     @client_context.require_multiple_mongoses
     @client_context.require_failCommand_fail_point
-    def test_retryable_reads_in_sharded_cluster_multiple_available(self):
+    def test_retryable_reads_are_retried_on_a_different_mongos_when_one_is_available(self):
         fail_command = {
             "configureFailPoint": "failCommand",
             "mode": {"times": 1},
-            "data": {
-                "failCommands": ["find"],
-                "closeConnection": True,
-                "appName": "retryableReadTest",
-            },
+            "data": {"failCommands": ["find"], "errorCode": 6},
         }
 
         mongos_clients = []
@@ -168,12 +164,11 @@ def test_retryable_reads_in_sharded_cluster_multiple_available(self):
         listener = OvertCommandListener()
         client = self.rs_or_single_client(
             client_context.mongos_seeds(),
-            appName="retryableReadTest",
             event_listeners=[listener],
             retryReads=True,
         )
 
-        with self.assertRaises(AutoReconnect):
+        with self.assertRaises(OperationFailure):
             client.t.t.find_one({})
 
         # Disable failpoints on each mongos
@@ -184,6 +179,43 @@ def test_retryable_reads_in_sharded_cluster_multiple_available(self):
         self.assertEqual(len(listener.failed_events), 2)
         self.assertEqual(len(listener.succeeded_events), 0)
 
+        #  Assert that both events occurred on different mongos.
+        assert listener.failed_events[0].connection_id != listener.failed_events[1].connection_id
+
+    @client_context.require_multiple_mongoses
+    @client_context.require_failCommand_fail_point
+    def test_retryable_reads_are_retried_on_the_same_mongos_when_no_others_are_available(self):
+        fail_command = {
+            "configureFailPoint": "failCommand",
+            "mode": {"times": 1},
+            "data": {"failCommands": ["find"], "errorCode": 6},
+        }
+
+        host = client_context.mongos_seeds().split(",")[0]
+        mongos_client = self.rs_or_single_client(host)
+        set_fail_point(mongos_client, fail_command)
+
+        listener = OvertCommandListener()
+        client = self.rs_or_single_client(
+            host,
+            directConnection=False,
+            event_listeners=[listener],
+            retryReads=True,
+        )
+
+        client.t.t.find_one({})
+
+        # Disable failpoint.
+        fail_command["mode"] = "off"
+        set_fail_point(mongos_client, fail_command)
+
+        # Assert that exactly one failed command event and one succeeded command event occurred.
+        self.assertEqual(len(listener.failed_events), 1)
+        self.assertEqual(len(listener.succeeded_events), 1)
+
+        #  Assert that both events occurred on the same mongos.
+        assert listener.succeeded_events[0].connection_id == listener.failed_events[0].connection_id
+
 
 if __name__ == "__main__":
     unittest.main()