PYTHON-3409 Retry flakey CSOT tests twice (#1046)

Author: ShaneHarvey

test/test_retryable_writes.py

@@ -471,6 +471,7 @@ def setUpClass(cls):
     def test_RetryableWriteError_error_label(self):
         listener = OvertCommandListener()
         client = rs_or_single_client(retryWrites=True, event_listeners=[listener])
+        self.addCleanup(client.close)
 
         # Ensure collection exists.
         client.pymongo_test.testcoll.insert_one({})

test/unified_format.py

@@ -25,6 +25,7 @@
 import re
 import sys
 import time
+import traceback
 import types
 from collections import abc
 from test import (
@@ -1580,6 +1581,25 @@ def verify_outcome(self, spec):
                 self.assertListEqual(sorted_expected_documents, actual_documents)
 
     def run_scenario(self, spec, uri=None):
+        if "csot" in self.id().lower():
+            # Retry CSOT tests up to 2 times to deal with flakey tests.
+            attempts = 3
+            for i in range(attempts):
+                try:
+                    return self._run_scenario(spec, uri)
+                except AssertionError:
+                    if i < attempts - 1:
+                        print(
+                            f"Retrying after attempt {i+1} of {self.id()} failed with:\n"
+                            f"{traceback.format_exc()}"
+                        )
+                        self.setUp()
+                        continue
+                    raise
+        else:
+            self._run_scenario(spec, uri)
+
+    def _run_scenario(self, spec, uri=None):
         # maybe skip test manually
         self.maybe_skip_test(spec)