From e55f829d27829961550174223a1649fb25e80568 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Mon, 5 May 2025 15:20:31 +0200 Subject: [PATCH] test: fix retry helpers currently causing flaky test failures Fix the retry helpers that are currently causing multiple tests to fail randomly. --- noxfile.py | 10 +++++++-- .../test_aborted_transaction.py | 19 +++++++++++++++++ tests/system/_helpers.py | 4 ++-- tests/system/test_dbapi.py | 21 ++++++++++++------- tests/system/test_session_api.py | 20 +++++++++--------- 5 files changed, 52 insertions(+), 22 deletions(-) diff --git a/noxfile.py b/noxfile.py index cb683afd7e..73ad757240 100644 --- a/noxfile.py +++ b/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) diff --git a/tests/mockserver_tests/test_aborted_transaction.py b/tests/mockserver_tests/test_aborted_transaction.py index 93eb42fe39..6a61dd4c73 100644 --- a/tests/mockserver_tests/test_aborted_transaction.py +++ b/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"]) diff --git a/tests/system/_helpers.py b/tests/system/_helpers.py index f37aefc2e5..1fc897b39c 100644 --- a/tests/system/_helpers.py +++ b/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): diff --git a/tests/system/test_dbapi.py b/tests/system/test_dbapi.py index a98f100bcc..6e4ced3c1b 100644 --- a/tests/system/test_dbapi.py +++ b/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() diff --git a/tests/system/test_session_api.py b/tests/system/test_session_api.py index 73b55b035d..21d7bccd44 100644 --- a/tests/system/test_session_api.py +++ b/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." )