PYTHON-5120 Reduce configureFailPoint duplication in tests (#2131)

Author: ShaneHarvey

test/__init__.py

@@ -933,17 +933,22 @@ 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, off=False):
+        cmd = {"configureFailPoint": "failCommand"}
+        cmd.update(command_args)
+        if off:
+            cmd["mode"] = "off"
+            cmd.pop("data", None)
+        client.admin.command(cmd)
+
     @contextmanager
     def fail_point(self, command_args):
-        cmd_on = SON([("configureFailPoint", "failCommand")])
-        cmd_on.update(command_args)
-        client_context.client.admin.command(cmd_on)
+        self.configure_fail_point(client_context.client, command_args)
         try:
             yield
         finally:
-            client_context.client.admin.command(
-                "configureFailPoint", cmd_on["configureFailPoint"], mode="off"
-            )
+            self.configure_fail_point(client_context.client, command_args, off=True)
 
     @contextmanager
     def fork(

test/asynchronous/__init__.py

@@ -935,17 +935,22 @@ 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, off=False):
+        cmd = {"configureFailPoint": "failCommand"}
+        cmd.update(command_args)
+        if off:
+            cmd["mode"] = "off"
+            cmd.pop("data", None)
+        await client.admin.command(cmd)
+
     @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)
+        await self.configure_fail_point(async_client_context.client, command_args)
         try:
             yield
         finally:
-            await async_client_context.client.admin.command(
-                "configureFailPoint", cmd_on["configureFailPoint"], mode="off"
-            )
+            await self.configure_fail_point(async_client_context.client, command_args, off=True)
 
     @contextmanager
     def fork(

test/asynchronous/test_connection_monitoring.py

@@ -211,15 +211,10 @@ def check_error(self, actual, expected):
         self.check_object(actual, expected)
         self.assertIn(message, str(actual))
 
-    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):
         if not async_client_context.supports_failCommand_fail_point:
             self.skipTest("failCommand fail point must be supported")
-        await self._set_fail_point(self.client, command_args)
+        await self.configure_fail_point(self.client, command_args)
 
     async def run_scenario(self, scenario_def, test):
         """Run a CMAP spec test."""

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):

test/asynchronous/unified_format.py

@@ -1008,12 +1008,8 @@ async def __set_fail_point(self, client, command_args):
         if not async_client_context.test_commands_enabled:
             self.skipTest("Test commands must be enabled")
 
-        cmd_on = SON([("configureFailPoint", "failCommand")])
-        cmd_on.update(command_args)
-        await client.admin.command(cmd_on)
-        self.addAsyncCleanup(
-            client.admin.command, "configureFailPoint", cmd_on["configureFailPoint"], mode="off"
-        )
+        await self.configure_fail_point(client, command_args)
+        self.addAsyncCleanup(self.configure_fail_point, client, command_args, off=True)
 
     async def _testOperation_failPoint(self, spec):
         await self.__set_fail_point(

test/asynchronous/utils_spec_runner.py

@@ -264,15 +264,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.
@@ -281,7 +276,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):

test/test_connection_monitoring.py

@@ -211,15 +211,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."""

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):

test/unified_format.py

@@ -999,12 +999,8 @@ def __set_fail_point(self, client, command_args):
         if not client_context.test_commands_enabled:
             self.skipTest("Test commands must be enabled")
 
-        cmd_on = SON([("configureFailPoint", "failCommand")])
-        cmd_on.update(command_args)
-        client.admin.command(cmd_on)
-        self.addCleanup(
-            client.admin.command, "configureFailPoint", cmd_on["configureFailPoint"], mode="off"
-        )
+        self.configure_fail_point(client, command_args)
+        self.addCleanup(self.configure_fail_point, client, command_args, off=True)
 
     def _testOperation_failPoint(self, spec):
         self.__set_fail_point(

test/utils_spec_runner.py

@@ -264,15 +264,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.
@@ -281,7 +276,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):