test: fix retry helpers currently causing flaky test failures

noxfile.py

@@ -51,6 +51,9 @@
     "pytest-cov",
     "pytest-asyncio",
 ]
+MOCK_SERVER_ADDITIONAL_DEPENDENCIES = [
+    "google-cloud-testutils",
+]
 UNIT_TEST_EXTERNAL_DEPENDENCIES: List[str] = []
 UNIT_TEST_LOCAL_DEPENDENCIES: List[str] = []
 UNIT_TEST_DEPENDENCIES: List[str] = []
@@ -242,8 +245,11 @@ def mockserver(session):
     constraints_path = str(
         CURRENT_DIRECTORY / "testing" / f"constraints-{session.python}.txt"
     )
-    # install_unittest_dependencies(session, "-c", constraints_path)
-    standard_deps = UNIT_TEST_STANDARD_DEPENDENCIES + UNIT_TEST_DEPENDENCIES
+    standard_deps = (
+        UNIT_TEST_STANDARD_DEPENDENCIES
+        + UNIT_TEST_DEPENDENCIES
+        + MOCK_SERVER_ADDITIONAL_DEPENDENCIES
+    )
     session.install(*standard_deps, "-c", constraints_path)
     session.install("-e", ".", "-c", constraints_path)
 

tests/mockserver_tests/test_aborted_transaction.py

@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import random
 
 from google.cloud.spanner_v1 import (
     BatchCreateSessionsRequest,
@@ -29,6 +30,12 @@
     add_update_count,
     add_single_result,
 )
+from google.api_core import exceptions
+from test_utils import retry
+
+retry_maybe_aborted_txn = retry.RetryErrors(
+    exceptions.Aborted, max_tries=5, delay=0, backoff=1
+)
 
 
 class TestAbortedTransaction(MockServerTestBase):
@@ -119,6 +126,18 @@ def test_batch_commit_aborted(self):
         # The transaction is aborted and retried.
         self.assertTrue(isinstance(requests[2], CommitRequest))
 
+    @retry_maybe_aborted_txn
+    def test_retry_helper(self):
+        # Randomly add an Aborted error for the Commit method on the mock server.
+        if random.random() < 0.5:
+            add_error(SpannerServicer.Commit.__name__, aborted_status())
+        session = self.database.session()
+        session.create()
+        transaction = session.transaction()
+        transaction.begin()
+        transaction.insert("my_table", ["col1, col2"], [{"col1": 1, "col2": "One"}])
+        transaction.commit()
+
 
 def _insert_mutations(transaction: Transaction):
     transaction.insert("my_table", ["col1", "col2"], ["value1", "value2"])

tests/system/_helpers.py

@@ -74,8 +74,8 @@
 retry_429_503 = retry.RetryErrors(
     exceptions.TooManyRequests, exceptions.ServiceUnavailable, 8
 )
-retry_mabye_aborted_txn = retry.RetryErrors(exceptions.ServerError, exceptions.Aborted)
-retry_mabye_conflict = retry.RetryErrors(exceptions.ServerError, exceptions.Conflict)
+retry_maybe_aborted_txn = retry.RetryErrors(exceptions.Aborted)
+retry_maybe_conflict = retry.RetryErrors(exceptions.Conflict)
 
 
 def _has_all_ddl(database):

tests/system/test_dbapi.py

@@ -763,12 +763,15 @@ def test_commit_abort_retry(self, dbapi_database):
         dbapi_database._method_abort_interceptor.set_method_to_abort(
             COMMIT_METHOD, self._conn
         )
-        # called 2 times
+        # called (at least) 2 times
         self._conn.commit()
         dbapi_database._method_abort_interceptor.reset()
-        assert method_count_interceptor._counts[COMMIT_METHOD] == 2
-        assert method_count_interceptor._counts[EXECUTE_BATCH_DML_METHOD] == 4
-        assert method_count_interceptor._counts[EXECUTE_STREAMING_SQL_METHOD] == 10
+        # Verify the number of calls.
+        # We don't know the exact number of calls, as Spanner could also
+        # abort the transaction.
+        assert method_count_interceptor._counts[COMMIT_METHOD] >= 2
+        assert method_count_interceptor._counts[EXECUTE_BATCH_DML_METHOD] >= 4
+        assert method_count_interceptor._counts[EXECUTE_STREAMING_SQL_METHOD] >= 10
 
         self._cursor.execute("SELECT * FROM contacts")
         got_rows = self._cursor.fetchall()
@@ -829,10 +832,12 @@ def test_execute_sql_abort_retry_multiple_times(self, dbapi_database):
         self._cursor.fetchmany(2)
         dbapi_database._method_abort_interceptor.reset()
         self._conn.commit()
-        # Check that all rpcs except commit should be called 3 times the original
-        assert method_count_interceptor._counts[COMMIT_METHOD] == 1
-        assert method_count_interceptor._counts[EXECUTE_BATCH_DML_METHOD] == 3
-        assert method_count_interceptor._counts[EXECUTE_STREAMING_SQL_METHOD] == 3
+        # Check that all RPCs except commit should be called at least 3 times
+        # We don't know the exact number of attempts, as the transaction could
+        # also be aborted by Spanner (and not only the test interceptor).
+        assert method_count_interceptor._counts[COMMIT_METHOD] >= 1
+        assert method_count_interceptor._counts[EXECUTE_BATCH_DML_METHOD] >= 3
+        assert method_count_interceptor._counts[EXECUTE_STREAMING_SQL_METHOD] >= 3
 
         self._cursor.execute("SELECT * FROM contacts")
         got_rows = self._cursor.fetchall()

tests/system/test_session_api.py

@@ -578,7 +578,7 @@ def test_batch_insert_w_commit_timestamp(sessions_database, not_postgres):
     assert not deleted
 
 
-@_helpers.retry_mabye_aborted_txn
+@_helpers.retry_maybe_aborted_txn
 def test_transaction_read_and_insert_then_rollback(
     sessions_database,
     ot_exporter,
@@ -687,7 +687,7 @@ def test_transaction_read_and_insert_then_rollback(
         )
 
 
-@_helpers.retry_mabye_conflict
+@_helpers.retry_maybe_conflict
 def test_transaction_read_and_insert_then_exception(sessions_database):
     class CustomException(Exception):
         pass
@@ -714,7 +714,7 @@ def _transaction_read_then_raise(transaction):
     assert rows == []
 
 
-@_helpers.retry_mabye_conflict
+@_helpers.retry_maybe_conflict
 def test_transaction_read_and_insert_or_update_then_commit(
     sessions_database,
     sessions_to_delete,
@@ -771,8 +771,8 @@ def _generate_insert_returning_statement(row, database_dialect):
     return f"INSERT INTO {table} ({column_list}) VALUES ({row_data}) {returning}"
 
 
-@_helpers.retry_mabye_conflict
-@_helpers.retry_mabye_aborted_txn
+@_helpers.retry_maybe_conflict
+@_helpers.retry_maybe_aborted_txn
 def test_transaction_execute_sql_w_dml_read_rollback(
     sessions_database,
     sessions_to_delete,
@@ -809,7 +809,7 @@ def test_transaction_execute_sql_w_dml_read_rollback(
     # [END spanner_test_dml_rollback_txn_not_committed]
 
 
-@_helpers.retry_mabye_conflict
+@_helpers.retry_maybe_conflict
 def test_transaction_execute_update_read_commit(sessions_database, sessions_to_delete):
     # [START spanner_test_dml_read_your_writes]
     sd = _sample_data
@@ -838,7 +838,7 @@ def test_transaction_execute_update_read_commit(sessions_database, sessions_to_d
     # [END spanner_test_dml_read_your_writes]
 
 
-@_helpers.retry_mabye_conflict
+@_helpers.retry_maybe_conflict
 def test_transaction_execute_update_then_insert_commit(
     sessions_database, sessions_to_delete
 ):
@@ -870,7 +870,7 @@ def test_transaction_execute_update_then_insert_commit(
     # [END spanner_test_dml_with_mutation]
 
 
-@_helpers.retry_mabye_conflict
+@_helpers.retry_maybe_conflict
 @pytest.mark.skipif(
     _helpers.USE_EMULATOR, reason="Emulator does not support DML Returning."
 )
@@ -901,7 +901,7 @@ def test_transaction_execute_sql_dml_returning(
     sd._check_rows_data(rows)
 
 
-@_helpers.retry_mabye_conflict
+@_helpers.retry_maybe_conflict
 @pytest.mark.skipif(
     _helpers.USE_EMULATOR, reason="Emulator does not support DML Returning."
 )
@@ -929,7 +929,7 @@ def test_transaction_execute_update_dml_returning(
     sd._check_rows_data(rows)
 
 
-@_helpers.retry_mabye_conflict
+@_helpers.retry_maybe_conflict
 @pytest.mark.skipif(
     _helpers.USE_EMULATOR, reason="Emulator does not support DML Returning."
 )