🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

Author: gcf-owl-bot[bot]

tests/mockserver_tests/mock_server_test_base.py

@@ -231,78 +231,108 @@ def database(self) -> Database:
             )
         return self._database
 
-    def assert_requests_sequence(self, requests, expected_types, transaction_type, allow_multiple_batch_create=True):
+    def assert_requests_sequence(
+        self,
+        requests,
+        expected_types,
+        transaction_type,
+        allow_multiple_batch_create=True,
+    ):
         """Assert that the requests sequence matches the expected types, accounting for multiplexed sessions and retries.
-        
+
         Args:
             requests: List of requests from spanner_service.requests
             expected_types: List of expected request types (excluding session creation requests)
             transaction_type: TransactionType enum value to check multiplexed session status
             allow_multiple_batch_create: If True, skip all leading BatchCreateSessionsRequest and one optional CreateSessionRequest
         """
-        from google.cloud.spanner_v1 import BatchCreateSessionsRequest, CreateSessionRequest
+        from google.cloud.spanner_v1 import (
+            BatchCreateSessionsRequest,
+            CreateSessionRequest,
+        )
+
         mux_enabled = is_multiplexed_enabled(transaction_type)
         idx = 0
         # Skip all leading BatchCreateSessionsRequest (for retries)
         if allow_multiple_batch_create:
-            while idx < len(requests) and isinstance(requests[idx], BatchCreateSessionsRequest):
+            while idx < len(requests) and isinstance(
+                requests[idx], BatchCreateSessionsRequest
+            ):
                 idx += 1
             # For multiplexed, optionally skip a CreateSessionRequest
-            if mux_enabled and idx < len(requests) and isinstance(requests[idx], CreateSessionRequest):
+            if (
+                mux_enabled
+                and idx < len(requests)
+                and isinstance(requests[idx], CreateSessionRequest)
+            ):
                 idx += 1
         else:
             if mux_enabled:
-                self.assertTrue(isinstance(requests[idx], BatchCreateSessionsRequest),
-                                f"Expected BatchCreateSessionsRequest at index {idx}, got {type(requests[idx])}")
+                self.assertTrue(
+                    isinstance(requests[idx], BatchCreateSessionsRequest),
+                    f"Expected BatchCreateSessionsRequest at index {idx}, got {type(requests[idx])}",
+                )
                 idx += 1
-                self.assertTrue(isinstance(requests[idx], CreateSessionRequest),
-                                f"Expected CreateSessionRequest at index {idx}, got {type(requests[idx])}")
+                self.assertTrue(
+                    isinstance(requests[idx], CreateSessionRequest),
+                    f"Expected CreateSessionRequest at index {idx}, got {type(requests[idx])}",
+                )
                 idx += 1
             else:
-                self.assertTrue(isinstance(requests[idx], BatchCreateSessionsRequest),
-                                f"Expected BatchCreateSessionsRequest at index {idx}, got {type(requests[idx])}")
+                self.assertTrue(
+                    isinstance(requests[idx], BatchCreateSessionsRequest),
+                    f"Expected BatchCreateSessionsRequest at index {idx}, got {type(requests[idx])}",
+                )
                 idx += 1
         # Check the rest of the expected request types
         for expected_type in expected_types:
-            self.assertTrue(isinstance(requests[idx], expected_type),
-                            f"Expected {expected_type} at index {idx}, got {type(requests[idx])}")
+            self.assertTrue(
+                isinstance(requests[idx], expected_type),
+                f"Expected {expected_type} at index {idx}, got {type(requests[idx])}",
+            )
             idx += 1
-        self.assertEqual(idx, len(requests),
-                         f"Expected {idx} requests, got {len(requests)}")
+        self.assertEqual(
+            idx, len(requests), f"Expected {idx} requests, got {len(requests)}"
+        )
 
     def adjust_request_id_sequence(self, expected_segments, requests, transaction_type):
         """Adjust expected request ID sequence numbers based on actual session creation requests.
-        
+
         Args:
             expected_segments: List of expected (method, (sequence_numbers)) tuples
             requests: List of actual requests from spanner_service.requests
             transaction_type: TransactionType enum value to check multiplexed session status
-            
+
         Returns:
             List of adjusted expected segments with corrected sequence numbers
         """
-        from google.cloud.spanner_v1 import BatchCreateSessionsRequest, CreateSessionRequest, ExecuteSqlRequest, BeginTransactionRequest
-        
+        from google.cloud.spanner_v1 import (
+            BatchCreateSessionsRequest,
+            CreateSessionRequest,
+            ExecuteSqlRequest,
+            BeginTransactionRequest,
+        )
+
         # Count session creation requests that come before the first non-session request
         session_requests_before = 0
         for req in requests:
             if isinstance(req, (BatchCreateSessionsRequest, CreateSessionRequest)):
                 session_requests_before += 1
             elif isinstance(req, (ExecuteSqlRequest, BeginTransactionRequest)):
                 break
-                
+
         # For multiplexed sessions, we expect 2 session requests (BatchCreateSessions + CreateSession)
         # For non-multiplexed, we expect 1 session request (BatchCreateSessions)
         mux_enabled = is_multiplexed_enabled(transaction_type)
         expected_session_requests = 2 if mux_enabled else 1
         extra_session_requests = session_requests_before - expected_session_requests
-                
+
         # Adjust sequence numbers based on extra session requests
         adjusted_segments = []
         for method, seq_nums in expected_segments:
             # Adjust the sequence number (5th element in the tuple)
             adjusted_seq_nums = list(seq_nums)
             adjusted_seq_nums[4] += extra_session_requests
             adjusted_segments.append((method, tuple(adjusted_seq_nums)))
-            
+
         return adjusted_segments

tests/mockserver_tests/test_aborted_transaction.py

@@ -50,7 +50,12 @@ def test_run_in_transaction_commit_aborted(self):
         requests = self.spanner_service.requests
         self.assert_requests_sequence(
             requests,
-            [BeginTransactionRequest, CommitRequest, BeginTransactionRequest, CommitRequest],
+            [
+                BeginTransactionRequest,
+                CommitRequest,
+                BeginTransactionRequest,
+                CommitRequest,
+            ],
             TransactionType.READ_WRITE,
         )
 

tests/mockserver_tests/test_basics.py

@@ -97,10 +97,20 @@ def test_dbapi_partitioned_dml(self):
         )
         # Find the first BeginTransactionRequest after session creation
         idx = 0
-        from google.cloud.spanner_v1 import BatchCreateSessionsRequest, CreateSessionRequest
-        while idx < len(requests) and isinstance(requests[idx], BatchCreateSessionsRequest):
+        from google.cloud.spanner_v1 import (
+            BatchCreateSessionsRequest,
+            CreateSessionRequest,
+        )
+
+        while idx < len(requests) and isinstance(
+            requests[idx], BatchCreateSessionsRequest
+        ):
             idx += 1
-        if is_multiplexed_enabled(TransactionType.PARTITIONED) and idx < len(requests) and isinstance(requests[idx], CreateSessionRequest):
+        if (
+            is_multiplexed_enabled(TransactionType.PARTITIONED)
+            and idx < len(requests)
+            and isinstance(requests[idx], CreateSessionRequest)
+        ):
             idx += 1
         begin_request: BeginTransactionRequest = requests[idx]
         self.assertEqual(

tests/mockserver_tests/test_request_id_header.py

@@ -59,7 +59,9 @@ def test_snapshot_execute_sql(self):
         CHANNEL_ID = self.database._channel_id
         got_stream_segments, got_unary_segments = self.canonicalize_request_id_headers()
         # Filter out CreateSessionRequest unary segments for comparison
-        filtered_unary_segments = [seg for seg in got_unary_segments if not seg[0].endswith("/CreateSession")]
+        filtered_unary_segments = [
+            seg for seg in got_unary_segments if not seg[0].endswith("/CreateSession")
+        ]
         want_unary_segments = [
             (
                 "/google.spanner.v1.Spanner/BatchCreateSessions",
@@ -76,7 +78,14 @@ def test_snapshot_execute_sql(self):
         want_stream_segments = [
             (
                 "/google.spanner.v1.Spanner/ExecuteStreamingSql",
-                (1, REQ_RAND_PROCESS_ID, NTH_CLIENT, CHANNEL_ID, 1 + session_requests_before, 1),
+                (
+                    1,
+                    REQ_RAND_PROCESS_ID,
+                    NTH_CLIENT,
+                    CHANNEL_ID,
+                    1 + session_requests_before,
+                    1,
+                ),
             )
         ]
         assert filtered_unary_segments == want_unary_segments
@@ -89,6 +98,7 @@ def test_snapshot_read_concurrent(self):
             rows = snapshot.execute_sql("select 1")
             for row in rows:
                 _ = row
+
         def select1():
             with db.snapshot() as snapshot:
                 rows = snapshot.execute_sql("select 1")
@@ -97,6 +107,7 @@ def select1():
                     self.assertEqual(1, row[0])
                     res_list.append(row)
                 self.assertEqual(1, len(res_list))
+
         n = 10
         threads = []
         for i in range(n):
@@ -110,7 +121,9 @@ def select1():
         # Allow for an extra request due to multiplexed session creation
         expected_min = 2 + n
         expected_max = expected_min + 1
-        assert expected_min <= len(requests) <= expected_max, f"Expected {expected_min} or {expected_max} requests, got {len(requests)}: {requests}"
+        assert (
+            expected_min <= len(requests) <= expected_max
+        ), f"Expected {expected_min} or {expected_max} requests, got {len(requests)}: {requests}"
         client_id = db._nth_client_id
         channel_id = db._channel_id
         got_stream_segments, got_unary_segments = self.canonicalize_request_id_headers()
@@ -132,7 +145,14 @@ def select1():
         want_stream_segments = [
             (
                 "/google.spanner.v1.Spanner/ExecuteStreamingSql",
-                (1, REQ_RAND_PROCESS_ID, client_id, channel_id, session_requests_before + i, 1),
+                (
+                    1,
+                    REQ_RAND_PROCESS_ID,
+                    client_id,
+                    channel_id,
+                    session_requests_before + i,
+                    1,
+                ),
             )
             for i in range(1, n + 2)
         ]
@@ -173,7 +193,9 @@ def test_database_execute_partitioned_dml_request_id(self):
         NTH_CLIENT = self.database._nth_client_id
         CHANNEL_ID = self.database._channel_id
         # Allow for extra unary segments due to session creation
-        filtered_unary_segments = [seg for seg in got_unary_segments if not seg[0].endswith("/CreateSession")]
+        filtered_unary_segments = [
+            seg for seg in got_unary_segments if not seg[0].endswith("/CreateSession")
+        ]
         # Find the actual sequence number for BeginTransaction
         begin_txn_seq = None
         for seg in filtered_unary_segments:

tests/mockserver_tests/test_tags.py

@@ -149,14 +149,25 @@ def test_select_read_write_transaction_with_transaction_tag(self):
         requests = self.spanner_service.requests
         self.assert_requests_sequence(
             requests,
-            [BeginTransactionRequest, ExecuteSqlRequest, ExecuteSqlRequest, CommitRequest],
+            [
+                BeginTransactionRequest,
+                ExecuteSqlRequest,
+                ExecuteSqlRequest,
+                CommitRequest,
+            ],
             TransactionType.READ_WRITE,
         )
         mux_enabled = is_multiplexed_enabled(TransactionType.READ_WRITE)
         tag_idx = 3 if mux_enabled else 2
-        self.assertEqual("my_transaction_tag", requests[tag_idx].request_options.transaction_tag)
-        self.assertEqual("my_transaction_tag", requests[tag_idx + 1].request_options.transaction_tag)
-        self.assertEqual("my_transaction_tag", requests[tag_idx + 2].request_options.transaction_tag)
+        self.assertEqual(
+            "my_transaction_tag", requests[tag_idx].request_options.transaction_tag
+        )
+        self.assertEqual(
+            "my_transaction_tag", requests[tag_idx + 1].request_options.transaction_tag
+        )
+        self.assertEqual(
+            "my_transaction_tag", requests[tag_idx + 2].request_options.transaction_tag
+        )
 
     def test_select_read_write_transaction_with_transaction_and_request_tag(self):
         connection = Connection(self.instance, self.database)
@@ -170,16 +181,27 @@ def test_select_read_write_transaction_with_transaction_and_request_tag(self):
         requests = self.spanner_service.requests
         self.assert_requests_sequence(
             requests,
-            [BeginTransactionRequest, ExecuteSqlRequest, ExecuteSqlRequest, CommitRequest],
+            [
+                BeginTransactionRequest,
+                ExecuteSqlRequest,
+                ExecuteSqlRequest,
+                CommitRequest,
+            ],
             TransactionType.READ_WRITE,
         )
         mux_enabled = is_multiplexed_enabled(TransactionType.READ_WRITE)
         tag_idx = 3 if mux_enabled else 2
-        self.assertEqual("my_transaction_tag", requests[tag_idx].request_options.transaction_tag)
+        self.assertEqual(
+            "my_transaction_tag", requests[tag_idx].request_options.transaction_tag
+        )
         self.assertEqual("my_tag1", requests[tag_idx].request_options.request_tag)
-        self.assertEqual("my_transaction_tag", requests[tag_idx + 1].request_options.transaction_tag)
+        self.assertEqual(
+            "my_transaction_tag", requests[tag_idx + 1].request_options.transaction_tag
+        )
         self.assertEqual("my_tag2", requests[tag_idx + 1].request_options.request_tag)
-        self.assertEqual("my_transaction_tag", requests[tag_idx + 2].request_options.transaction_tag)
+        self.assertEqual(
+            "my_transaction_tag", requests[tag_idx + 2].request_options.transaction_tag
+        )
 
     def test_request_tag_is_cleared(self):
         connection = Connection(self.instance, self.database)