Update tests

Author: odeke-em

google/cloud/spanner_v1/database.py

@@ -749,14 +749,15 @@ def execute_partitioned_dml(
                 _metadata_with_leader_aware_routing(self._route_to_leader_enabled)
             )
 
-        # Attempt will be incremented inside _restart_on_unavailable.
         begin_txn_nth_request = self._next_nth_request
-        begin_txn_attempt = AtomicCounter(1)
+        begin_txn_attempt = AtomicCounter(0)
         partial_nth_request = self._next_nth_request
+        # partial_attempt will be incremented inside _restart_on_unavailable.
         partial_attempt = AtomicCounter(0)
 
         def execute_pdml():
             with SessionCheckout(self._pool) as session:
+                begin_txn_attempt.increment()
                 txn = api.begin_transaction(
                     session=session.name,
                     options=txn_options,
@@ -792,7 +793,6 @@ def wrapped_method(*args, **kwargs):
                     request=request,
                     transaction_selector=txn_selector,
                     observability_options=self.observability_options,
-                    attempt=begin_txn_attempt,
                 )
 
                 result_set = StreamedResultSet(iterator)

google/cloud/spanner_v1/instance.py

@@ -501,7 +501,6 @@ def database(
                 proto_descriptors=proto_descriptors,
             )
         else:
-            print("enabled interceptors")
             return TestDatabase(
                 database_id,
                 self,

google/cloud/spanner_v1/pool.py

@@ -243,7 +243,6 @@ def bind(self, database):
             "CloudSpanner.FixedPool.BatchCreateSessions",
             observability_options=observability_options,
         ) as span:
-            attempt = 1
             returned_session_count = 0
             while not self._sessions.full():
                 request.session_count = requested_session_count - self._sessions.qsize()
@@ -253,7 +252,7 @@ def bind(self, database):
                     span_event_attributes,
                 )
                 all_metadata = database.metadata_with_request_id(
-                    database._next_nth_request, attempt, metadata
+                    database._next_nth_request, 1, metadata
                 )
                 resp = api.batch_create_sessions(
                     request=request,

google/cloud/spanner_v1/snapshot.py

@@ -35,7 +35,6 @@
     _merge_query_options,
     _metadata_with_prefix,
     _metadata_with_leader_aware_routing,
-    _metadata_with_request_id,
     _retry,
     _check_rst_stream_error,
     _SessionWrapper,
@@ -60,7 +59,6 @@ def _restart_on_unavailable(
     transaction=None,
     transaction_selector=None,
     observability_options=None,
-    attempt=0,
 ):
     """Restart iteration after :exc:`.ServiceUnavailable`.
 
@@ -95,7 +93,6 @@ def _restart_on_unavailable(
     ):
         iterator = method(request=request)
     while True:
-        attempt += 1
         try:
             for item in iterator:
                 item_buffer.append(item)

google/cloud/spanner_v1/testing/interceptors.py

@@ -91,13 +91,6 @@ def intercept(self, method, request_or_iterator, call_details):
 
         response_or_iterator = method(request_or_iterator, call_details)
         streaming = getattr(response_or_iterator, "__iter__", None) is not None
-        print(
-            "intercept got",
-            x_goog_request_id,
-            call_details.method,
-            "streaming",
-            streaming,
-        )
         with self.__lock:
             if streaming:
                 self._stream_req_segments.append(

google/cloud/spanner_v1/testing/mock_spanner.py

@@ -22,8 +22,6 @@
 from google.cloud.spanner_v1 import (
     TransactionOptions,
     ResultSetMetadata,
-    ExecuteSqlRequest,
-    ExecuteBatchDmlRequest,
 )
 from google.cloud.spanner_v1.testing.mock_database_admin import DatabaseAdminServicer
 import google.cloud.spanner_v1.testing.spanner_database_admin_pb2_grpc as database_admin_grpc
@@ -186,9 +184,7 @@ def BeginTransaction(self, request, context):
         self._requests.append(request)
         return self.__create_transaction(request.session, request.options)
 
-    def __maybe_create_transaction(
-        self, request: ExecuteSqlRequest | ExecuteBatchDmlRequest
-    ):
+    def __maybe_create_transaction(self, request):
         started_transaction = None
         if not request.transaction.begin == TransactionOptions():
             started_transaction = self.__create_transaction(

google/cloud/spanner_v1/transaction.py

@@ -22,7 +22,6 @@
     _merge_query_options,
     _metadata_with_prefix,
     _metadata_with_leader_aware_routing,
-    _metadata_with_request_id,
     _retry,
     _check_rst_stream_error,
 )

tests/unit/test_atomic_counter.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import time
 import random
 import threading
 import unittest

tests/unit/test_database.py

@@ -3258,7 +3258,6 @@ def _nth_client_id(self):
         return self._instance._client._nth_client_id
 
     def metadata_with_request_id(self, nth_request, nth_attempt, prior_metadata=[]):
-        client_id = self._nth_client_id
         return _metadata_with_request_id(
             self._nth_client_id,
             self._channel_id,

tests/unit/test_pool.py

@@ -1250,7 +1250,6 @@ def _nth_client_id(self):
         return self.NTH_CLIENT_ID.increment()
 
     def metadata_with_request_id(self, nth_request, nth_attempt, prior_metadata=[]):
-        client_id = self._nth_client_id
         return _metadata_with_request_id(
             self._nth_client_id,
             self._channel_id,