test: fix retry helpers currently causing flaky test failures

Fix the retry helpers that are currently causing multiple tests to fail randomly.

Author: olavloite

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_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."
 )