Add other retryable reads to test

Author: NoahStapp

test/asynchronous/test_retryable_reads.py

@@ -220,35 +220,43 @@ async def test_retryable_reads_are_retried_on_the_same_mongos_when_no_others_are
 
     @async_client_context.require_failCommand_fail_point
     async def test_retryable_reads_are_retried_on_the_same_implicit_session(self):
-        fail_command = {
-            "configureFailPoint": "failCommand",
-            "mode": {"times": 1},
-            "data": {"failCommands": ["count"], "errorCode": 6},
-        }
-
         listener = OvertCommandListener()
         client = await self.async_rs_or_single_client(
             directConnection=False,
             event_listeners=[listener],
             retryReads=True,
         )
 
-        await async_set_fail_point(client, fail_command)
-
-        await client.t.t.estimated_document_count()
+        commands = [
+            ("aggregate", lambda: client.t.t.count_documents({})),
+            ("aggregate", lambda: client.t.t.aggregate([{"$match": {}}])),
+            ("count", lambda: client.t.t.estimated_document_count()),
+            ("distinct", lambda: client.t.t.distinct("x")),
+            ("find", lambda: client.t.t.find_one({})),
+            ("listDatabases", lambda: client.list_databases()),
+            ("listCollections", lambda: client.t.list_collections()),
+            ("listIndexes", lambda: client.t.t.list_indexes()),
+        ]
 
-        # Disable failpoint.
-        fail_command["mode"] = "off"
-        await async_set_fail_point(client, fail_command)
+        for command_name, operation in commands:
+            listener.reset()
+            fail_command = {
+                "configureFailPoint": "failCommand",
+                "mode": {"times": 1},
+                "data": {"failCommands": [command_name], "errorCode": 6},
+            }
 
-        #  Assert that both events occurred on the same session.
-        lsids = [
-            event.command["lsid"]
-            for event in listener.started_events
-            if event.command_name == "count"
-        ]
-        self.assertEqual(len(lsids), 2)
-        self.assertEqual(lsids[0], lsids[1])
+            async with self.fail_point(fail_command):
+                await operation()
+
+            #  Assert that both events occurred on the same session.
+            lsids = [
+                event.command["lsid"]
+                for event in listener.started_events
+                if event.command_name == command_name
+            ]
+            self.assertEqual(len(lsids), 2)
+            self.assertEqual(lsids[0], lsids[1])
 
 
 if __name__ == "__main__":

test/test_retryable_reads.py

@@ -218,35 +218,43 @@ def test_retryable_reads_are_retried_on_the_same_mongos_when_no_others_are_avail
 
     @client_context.require_failCommand_fail_point
     def test_retryable_reads_are_retried_on_the_same_implicit_session(self):
-        fail_command = {
-            "configureFailPoint": "failCommand",
-            "mode": {"times": 1},
-            "data": {"failCommands": ["count"], "errorCode": 6},
-        }
-
         listener = OvertCommandListener()
         client = self.rs_or_single_client(
             directConnection=False,
             event_listeners=[listener],
             retryReads=True,
         )
 
-        set_fail_point(client, fail_command)
-
-        client.t.t.estimated_document_count()
+        commands = [
+            ("aggregate", lambda: client.t.t.count_documents({})),
+            ("aggregate", lambda: client.t.t.aggregate([{"$match": {}}])),
+            ("count", lambda: client.t.t.estimated_document_count()),
+            ("distinct", lambda: client.t.t.distinct("x")),
+            ("find", lambda: client.t.t.find_one({})),
+            ("listDatabases", lambda: client.list_databases()),
+            ("listCollections", lambda: client.t.list_collections()),
+            ("listIndexes", lambda: client.t.t.list_indexes()),
+        ]
 
-        # Disable failpoint.
-        fail_command["mode"] = "off"
-        set_fail_point(client, fail_command)
+        for command_name, operation in commands:
+            listener.reset()
+            fail_command = {
+                "configureFailPoint": "failCommand",
+                "mode": {"times": 1},
+                "data": {"failCommands": [command_name], "errorCode": 6},
+            }
 
-        #  Assert that both events occurred on the same session.
-        lsids = [
-            event.command["lsid"]
-            for event in listener.started_events
-            if event.command_name == "count"
-        ]
-        self.assertEqual(len(lsids), 2)
-        self.assertEqual(lsids[0], lsids[1])
+            with self.fail_point(fail_command):
+                operation()
+
+            #  Assert that both events occurred on the same session.
+            lsids = [
+                event.command["lsid"]
+                for event in listener.started_events
+                if event.command_name == command_name
+            ]
+            self.assertEqual(len(lsids), 2)
+            self.assertEqual(lsids[0], lsids[1])
 
 
 if __name__ == "__main__":