PYTHON-5071 Fix configureFailPoint retry

ShaneHarvey · ShaneHarvey · commit 758483165e4c · 2025-01-31T13:52:48.000-08:00
diff --git a/test/__init__.py b/test/__init__.py
@@ -933,25 +933,29 @@ def assertEqualCommand(self, expected, actual, msg=None):
     def assertEqualReply(self, expected, actual, msg=None):
         self.assertEqual(sanitize_reply(expected), sanitize_reply(actual), msg)
 
+    @staticmethod
+    def configure_fail_point(client, command_args):
+        cmd = {"configureFailPoint": "failCommand"}
+        cmd.update(command_args)
+        try:
+            client.admin.command(cmd)
+            return
+        except pymongo.errors.ConnectionFailure:
+            # Workaround PyPy bug described in PYTHON-5011.
+            if not _IS_SYNC and "PyPy" in sys.version:
+                client.admin.command(cmd)
+                return
+            raise
+
     @contextmanager
     def fail_point(self, command_args):
-        cmd_on = SON([("configureFailPoint", "failCommand")])
-        cmd_on.update(command_args)
-        client_context.client.admin.command(cmd_on)
+        name = command_args.get("configureFailPoint", "failCommand")
+        self.configure_fail_point(client_context.client, command_args)
         try:
             yield
         finally:
-            try:
-                client_context.client.admin.command(
-                    "configureFailPoint", cmd_on["configureFailPoint"], mode="off"
-                )
-            except pymongo.errors.ConnectionFailure:
-                # Workaround PyPy bug described in PYTHON-5011.
-                if not _IS_SYNC and "PyPy" in sys.version:
-                    client_context.client.admin.command(
-                        "configureFailPoint", cmd_on["configureFailPoint"], mode="off"
-                    )
-                raise
+            cmd_off = {"configureFailPoint": name, "mode": "off"}
+            self.configure_fail_point(client_context.client, cmd_off)
 
     @contextmanager
     def fork(
diff --git a/test/asynchronous/__init__.py b/test/asynchronous/__init__.py
@@ -935,25 +935,29 @@ def assertEqualCommand(self, expected, actual, msg=None):
     def assertEqualReply(self, expected, actual, msg=None):
         self.assertEqual(sanitize_reply(expected), sanitize_reply(actual), msg)
 
+    @staticmethod
+    async def configure_fail_point(client, command_args):
+        cmd = {"configureFailPoint": "failCommand"}
+        cmd.update(command_args)
+        try:
+            await client.admin.command(cmd)
+            return
+        except pymongo.errors.ConnectionFailure:
+            # Workaround PyPy bug described in PYTHON-5011.
+            if not _IS_SYNC and "PyPy" in sys.version:
+                await client.admin.command(cmd)
+                return
+            raise
+
     @asynccontextmanager
     async def fail_point(self, command_args):
-        cmd_on = SON([("configureFailPoint", "failCommand")])
-        cmd_on.update(command_args)
-        await async_client_context.client.admin.command(cmd_on)
+        name = command_args.get("configureFailPoint", "failCommand")
+        await self.configure_fail_point(async_client_context.client, command_args)
         try:
             yield
         finally:
-            try:
-                await async_client_context.client.admin.command(
-                    "configureFailPoint", cmd_on["configureFailPoint"], mode="off"
-                )
-            except pymongo.errors.ConnectionFailure:
-                # Workaround PyPy bug described in PYTHON-5011.
-                if not _IS_SYNC and "PyPy" in sys.version:
-                    await async_client_context.client.admin.command(
-                        "configureFailPoint", cmd_on["configureFailPoint"], mode="off"
-                    )
-                raise
+            cmd_off = {"configureFailPoint": name, "mode": "off"}
+            await self.configure_fail_point(async_client_context.client, cmd_off)
 
     @contextmanager
     def fork(
diff --git a/test/asynchronous/test_transactions.py b/test/asynchronous/test_transactions.py
@@ -410,15 +410,10 @@ async def asyncSetUp(self) -> None:
             for address in async_client_context.mongoses:
                 self.mongos_clients.append(await self.async_single_client("{}:{}".format(*address)))
 
-    async def _set_fail_point(self, client, command_args):
-        cmd = {"configureFailPoint": "failCommand"}
-        cmd.update(command_args)
-        await client.admin.command(cmd)
-
     async def set_fail_point(self, command_args):
         clients = self.mongos_clients if self.mongos_clients else [self.client]
         for client in clients:
-            await self._set_fail_point(client, command_args)
+            await self.configure_fail_point(client, command_args)
 
     @async_client_context.require_transactions
     async def test_callback_raises_custom_error(self):
diff --git a/test/asynchronous/utils_spec_runner.py b/test/asynchronous/utils_spec_runner.py
@@ -265,15 +265,10 @@ async def asyncSetUp(self) -> None:
     async def asyncTearDown(self) -> None:
         self.knobs.disable()
 
-    async def _set_fail_point(self, client, command_args):
-        cmd = SON([("configureFailPoint", "failCommand")])
-        cmd.update(command_args)
-        await client.admin.command(cmd)
-
     async def set_fail_point(self, command_args):
         clients = self.mongos_clients if self.mongos_clients else [self.client]
         for client in clients:
-            await self._set_fail_point(client, command_args)
+            await self.configure_fail_point(client, command_args)
 
     async def targeted_fail_point(self, session, fail_point):
         """Run the targetedFailPoint test operation.
@@ -282,7 +277,7 @@ async def targeted_fail_point(self, session, fail_point):
         """
         clients = {c.address: c for c in self.mongos_clients}
         client = clients[session._pinned_address]
-        await self._set_fail_point(client, fail_point)
+        await self.configure_fail_point(client, fail_point)
         self.addAsyncCleanup(self.set_fail_point, {"mode": "off"})
 
     def assert_session_pinned(self, session):
diff --git a/test/test_connection_monitoring.py b/test/test_connection_monitoring.py
@@ -204,15 +204,10 @@ def check_error(self, actual, expected):
         self.check_object(actual, expected)
         self.assertIn(message, str(actual))
 
-    def _set_fail_point(self, client, command_args):
-        cmd = SON([("configureFailPoint", "failCommand")])
-        cmd.update(command_args)
-        client.admin.command(cmd)
-
     def set_fail_point(self, command_args):
         if not client_context.supports_failCommand_fail_point:
             self.skipTest("failCommand fail point must be supported")
-        self._set_fail_point(self.client, command_args)
+        self.configure_fail_point(self.client, command_args)
 
     def run_scenario(self, scenario_def, test):
         """Run a CMAP spec test."""
diff --git a/test/test_transactions.py b/test/test_transactions.py
@@ -402,15 +402,10 @@ def setUp(self) -> None:
             for address in client_context.mongoses:
                 self.mongos_clients.append(self.single_client("{}:{}".format(*address)))
 
-    def _set_fail_point(self, client, command_args):
-        cmd = {"configureFailPoint": "failCommand"}
-        cmd.update(command_args)
-        client.admin.command(cmd)
-
     def set_fail_point(self, command_args):
         clients = self.mongos_clients if self.mongos_clients else [self.client]
         for client in clients:
-            self._set_fail_point(client, command_args)
+            self.configure_fail_point(client, command_args)
 
     @client_context.require_transactions
     def test_callback_raises_custom_error(self):
diff --git a/test/utils_spec_runner.py b/test/utils_spec_runner.py
@@ -265,15 +265,10 @@ def setUp(self) -> None:
     def tearDown(self) -> None:
         self.knobs.disable()
 
-    def _set_fail_point(self, client, command_args):
-        cmd = SON([("configureFailPoint", "failCommand")])
-        cmd.update(command_args)
-        client.admin.command(cmd)
-
     def set_fail_point(self, command_args):
         clients = self.mongos_clients if self.mongos_clients else [self.client]
         for client in clients:
-            self._set_fail_point(client, command_args)
+            self.configure_fail_point(client, command_args)
 
     def targeted_fail_point(self, session, fail_point):
         """Run the targetedFailPoint test operation.
@@ -282,7 +277,7 @@ def targeted_fail_point(self, session, fail_point):
         """
         clients = {c.address: c for c in self.mongos_clients}
         client = clients[session._pinned_address]
-        self._set_fail_point(client, fail_point)
+        self.configure_fail_point(client, fail_point)
         self.addCleanup(self.set_fail_point, {"mode": "off"})
 
     def assert_session_pinned(self, session):
