increase retries for csot

Author: blink1073

test/asynchronous/unified_format.py

@@ -1369,7 +1369,6 @@ async def run_scenario(self, spec, uri=None):
         await self.kill_all_sessions()
 
         # Handle flaky tests.
-        func = functools.partial(self._run_scenario, spec, uri)
         flaky_tests = [
             # PYTHON-5170
             ".*test_discovery_and_monitoring.TestUnifiedInterruptInUsePoolClear.*",
@@ -1379,10 +1378,14 @@ async def run_scenario(self, spec, uri=None):
         ]
         for flaky_test in flaky_tests:
             if re.match(flaky_test, self.id()) is not None:
-                decorator = flaky(reset_func=self.asyncSetUp, func_name=self.id())
-                await decorator(func)()
+                func_name = self.id()
+                options = dict(reset_func=self.asyncSetUp, func_name=func_name)
+                if "csot" in func_name:
+                    options["max_runs"] = 3
+                decorator = flaky(**options)
+                await decorator(self._run_scenario)(spec, uri)
                 return
-        await func()
+        await self._run_scenario(spec, uri)
 
     async def _run_scenario(self, spec, uri=None):
         # maybe skip test manually

test/unified_format.py

@@ -1356,7 +1356,6 @@ def run_scenario(self, spec, uri=None):
         self.kill_all_sessions()
 
         # Handle flaky tests.
-        func = functools.partial(self._run_scenario, spec, uri)
         flaky_tests = [
             # PYTHON-5170
             ".*test_discovery_and_monitoring.TestUnifiedInterruptInUsePoolClear.*",
@@ -1366,10 +1365,14 @@ def run_scenario(self, spec, uri=None):
         ]
         for flaky_test in flaky_tests:
             if re.match(flaky_test, self.id()) is not None:
-                decorator = flaky(reset_func=self.setUp, func_name=self.id())
-                decorator(func)()
+                func_name = self.id()
+                options = dict(reset_func=self.setUp, func_name=func_name)
+                if "csot" in func_name:
+                    options["max_runs"] = 3
+                decorator = flaky(**options)
+                decorator(self._run_scenario)(spec, uri)
                 return
-        func()
+        self._run_scenario(spec, uri)
 
     def _run_scenario(self, spec, uri=None):
         # maybe skip test manually