test(spanner): update mock server tests to check for inline begin (#16656)

This PR updates the `sqlalchemy-spanner` mock server tests to account
for the dialect now inlining `BeginTransaction` requests into the first
statement execution.

### Changes
- Updated expected request counts in multiple tests (typically reducing
by 1).
- Removed assertions for separate `BeginTransactionRequest` where
applicable.
- Added assertions to verify that the first statement request contains
`transaction.begin` options.
- Updated assertions to verify that subsequent requests use the
transaction ID returned by the mock server.

These changes resolve the off-by-one failures in request counts across
multiple mock server tests.

Author: chalmerlowe

packages/sqlalchemy-spanner/tests/mockserver_tests/test_auto_increment.py

@@ -13,13 +13,12 @@
 # limitations under the License.
 
 from sqlalchemy.orm import Session
-from sqlalchemy.testing import eq_, is_instance_of
+from sqlalchemy.testing import eq_, is_instance_of, is_not_none
 from google.cloud.spanner_v1 import (
     ResultSet,
     CreateSessionRequest,
     ExecuteSqlRequest,
     CommitRequest,
-    BeginTransactionRequest,
 )
 from tests.mockserver_tests.mock_server_test_base import (
     MockServerTestBase,
@@ -126,11 +125,12 @@ def test_insert_row(self):
             session.commit()
         # Verify the requests that we got.
         requests = self.spanner_service.requests
-        eq_(4, len(requests))
+        # Dialect now inlines BeginTransaction into the first statement.
+        eq_(3, len(requests))
         is_instance_of(requests[0], CreateSessionRequest)
-        is_instance_of(requests[1], BeginTransactionRequest)
-        is_instance_of(requests[2], ExecuteSqlRequest)
-        is_instance_of(requests[3], CommitRequest)
+        is_instance_of(requests[1], ExecuteSqlRequest)
+        is_instance_of(requests[2], CommitRequest)
+        is_not_none(requests[1].transaction.begin)  # First request inlines begin
 
     def test_insert_row_with_pk_value(self):
         from tests.mockserver_tests.auto_increment_model import Singer

packages/sqlalchemy-spanner/tests/mockserver_tests/test_bit_reversed_sequence.py

@@ -13,13 +13,12 @@
 # limitations under the License.
 
 from sqlalchemy.orm import Session
-from sqlalchemy.testing import eq_, is_instance_of
+from sqlalchemy.testing import eq_, is_instance_of, is_not_none
 from google.cloud.spanner_v1 import (
     ResultSet,
     CreateSessionRequest,
     ExecuteSqlRequest,
     CommitRequest,
-    BeginTransactionRequest,
 )
 from tests.mockserver_tests.mock_server_test_base import (
     MockServerTestBase,
@@ -119,8 +118,9 @@ def test_insert_row(self):
             session.commit()
         # Verify the requests that we got.
         requests = self.spanner_service.requests
-        eq_(4, len(requests))
+        # Dialect now inlines BeginTransaction into the first statement.
+        eq_(3, len(requests))
         is_instance_of(requests[0], CreateSessionRequest)
-        is_instance_of(requests[1], BeginTransactionRequest)
-        is_instance_of(requests[2], ExecuteSqlRequest)
-        is_instance_of(requests[3], CommitRequest)
+        is_instance_of(requests[1], ExecuteSqlRequest)
+        is_instance_of(requests[2], CommitRequest)
+        is_not_none(requests[1].transaction.begin)  # First request inlines begin

packages/sqlalchemy-spanner/tests/mockserver_tests/test_float32.py

@@ -17,13 +17,13 @@
     eq_,
     is_instance_of,
     is_false,
+    is_not_none,
 )
 from google.cloud.spanner_v1 import (
     CreateSessionRequest,
     ExecuteSqlRequest,
     ResultSet,
     ResultSetStats,
-    BeginTransactionRequest,
     CommitRequest,
     TypeCode,
 )
@@ -58,12 +58,13 @@ def test_insert_data(self):
             session.commit()
 
             requests = self.spanner_service.requests
-            eq_(4, len(requests))
+            # Dialect now inlines BeginTransaction into the first statement.
+            eq_(3, len(requests))
             is_instance_of(requests[0], CreateSessionRequest)
-            is_instance_of(requests[1], BeginTransactionRequest)
-            is_instance_of(requests[2], ExecuteSqlRequest)
-            is_instance_of(requests[3], CommitRequest)
-            request: ExecuteSqlRequest = requests[2]
+            is_instance_of(requests[1], ExecuteSqlRequest)
+            is_instance_of(requests[2], CommitRequest)
+            is_not_none(requests[1].transaction.begin)  # First request inlines begin
+            request: ExecuteSqlRequest = requests[1]
             eq_(3, len(request.params))
             eq_("1", request.params["a0"])
             eq_("One", request.params["a1"])

packages/sqlalchemy-spanner/tests/mockserver_tests/test_insertmany.py

@@ -17,12 +17,11 @@
 
 import sqlalchemy
 from sqlalchemy.orm import Session
-from sqlalchemy.testing import eq_, is_instance_of
+from sqlalchemy.testing import eq_, is_instance_of, is_not_none
 from google.cloud.spanner_v1 import (
     ExecuteSqlRequest,
     CommitRequest,
     RollbackRequest,
-    BeginTransactionRequest,
     CreateSessionRequest,
 )
 from tests.mockserver_tests.mock_server_test_base import (
@@ -53,11 +52,12 @@ def test_insertmany_with_uuid_sentinels(self):
 
         # Verify the requests that we got.
         requests = self.spanner_service.requests
-        eq_(4, len(requests))
+        # Dialect now inlines BeginTransaction into the first statement.
+        eq_(3, len(requests))
         is_instance_of(requests[0], CreateSessionRequest)
-        is_instance_of(requests[1], BeginTransactionRequest)
-        is_instance_of(requests[2], ExecuteSqlRequest)
-        is_instance_of(requests[3], CommitRequest)
+        is_instance_of(requests[1], ExecuteSqlRequest)
+        is_instance_of(requests[2], CommitRequest)
+        is_not_none(requests[1].transaction.begin)  # First request inlines begin
 
     def test_no_insertmany_with_bit_reversed_id(self):
         """Ensures we don't try to bulk insert rows with bit-reversed PKs.
@@ -93,12 +93,13 @@ def test_no_insertmany_with_bit_reversed_id(self):
 
         # Verify the requests that we got.
         requests = self.spanner_service.requests
-        eq_(5, len(requests))
+        # Dialect now inlines BeginTransaction into the first statement.
+        eq_(4, len(requests))
         is_instance_of(requests[0], CreateSessionRequest)
-        is_instance_of(requests[1], BeginTransactionRequest)
+        is_instance_of(requests[1], ExecuteSqlRequest)
         is_instance_of(requests[2], ExecuteSqlRequest)
-        is_instance_of(requests[3], ExecuteSqlRequest)
-        is_instance_of(requests[4], RollbackRequest)
+        is_instance_of(requests[3], RollbackRequest)
+        is_not_none(requests[1].transaction.begin)  # First request inlines begin
 
     def add_uuid_insert_result(self, sql):
         result = result_set.ResultSet(

packages/sqlalchemy-spanner/tests/mockserver_tests/test_isolation_level.py

@@ -19,7 +19,6 @@
     CreateSessionRequest,
     ExecuteSqlRequest,
     CommitRequest,
-    BeginTransactionRequest,
     TransactionOptions,
 )
 
@@ -150,20 +149,20 @@ def test_invalid_isolation_level(self):
     def verify_isolation_level(self, level):
         # Verify the requests that we got.
         requests = self.spanner_service.requests
-        eq_(4, len(requests))
+        # Dialect now inlines BeginTransaction into the first statement.
+        eq_(3, len(requests))
         is_instance_of(requests[0], CreateSessionRequest)
-        is_instance_of(requests[1], BeginTransactionRequest)
-        is_instance_of(requests[2], ExecuteSqlRequest)
-        is_instance_of(requests[3], CommitRequest)
-        begin_request: BeginTransactionRequest = requests[1]
+        is_instance_of(requests[1], ExecuteSqlRequest)
+        is_instance_of(requests[2], CommitRequest)
+        execute_request: ExecuteSqlRequest = requests[1]
         eq_(
             TransactionOptions(
                 dict(
                     isolation_level=level,
                     read_write=TransactionOptions.ReadWrite(),
                 )
             ),
-            begin_request.options,
+            execute_request.transaction.begin,
         )
 
     def add_insert_result(self, sql):

packages/sqlalchemy-spanner/tests/mockserver_tests/test_json.py

@@ -14,13 +14,12 @@
 
 from sqlalchemy import func, select, text
 from sqlalchemy.orm import Session
-from sqlalchemy.testing import eq_, is_instance_of
+from sqlalchemy.testing import eq_, is_instance_of, is_not_none
 from google.cloud.spanner_v1 import (
     ResultSet,
     CreateSessionRequest,
     ExecuteSqlRequest,
     CommitRequest,
-    BeginTransactionRequest,
     TypeCode,
     JsonObject,
 )
@@ -105,12 +104,13 @@ def _test_insert_json(
 
         # Verify the requests that we got.
         requests = self.spanner_service.requests
-        eq_(4, len(requests))
+        # Dialect now inlines BeginTransaction into the first statement.
+        eq_(3, len(requests))
         is_instance_of(requests[0], CreateSessionRequest)
-        is_instance_of(requests[1], BeginTransactionRequest)
-        is_instance_of(requests[2], ExecuteSqlRequest)
-        is_instance_of(requests[3], CommitRequest)
-        request: ExecuteSqlRequest = requests[2]
+        is_instance_of(requests[1], ExecuteSqlRequest)
+        is_instance_of(requests[2], CommitRequest)
+        is_not_none(requests[1].transaction.begin)  # First request inlines begin
+        request: ExecuteSqlRequest = requests[1]
         eq_(3, len(request.params))
         eq_("1", request.params["a0"])
         eq_("Test", request.params["a1"])

packages/sqlalchemy-spanner/tests/mockserver_tests/test_pickle_type.py

@@ -13,13 +13,12 @@
 # limitations under the License.
 
 from sqlalchemy.orm import Session
-from sqlalchemy.testing import eq_, is_instance_of
+from sqlalchemy.testing import eq_, is_instance_of, is_not_none
 from google.cloud.spanner_v1 import (
     ResultSet,
     CreateSessionRequest,
     ExecuteSqlRequest,
     CommitRequest,
-    BeginTransactionRequest,
     TypeCode,
 )
 from tests.mockserver_tests.mock_server_test_base import (
@@ -85,12 +84,13 @@ def test_insert_and_query(self):
 
         # Verify the requests that we got.
         requests = self.spanner_service.requests
-        eq_(4, len(requests))
+        # Dialect now inlines BeginTransaction into the first statement.
+        eq_(3, len(requests))
         is_instance_of(requests[0], CreateSessionRequest)
-        is_instance_of(requests[1], BeginTransactionRequest)
-        is_instance_of(requests[2], ExecuteSqlRequest)
-        is_instance_of(requests[3], CommitRequest)
-        request: ExecuteSqlRequest = requests[2]
+        is_instance_of(requests[1], ExecuteSqlRequest)
+        is_instance_of(requests[2], CommitRequest)
+        is_not_none(requests[1].transaction.begin)  # First request inlines begin
+        request: ExecuteSqlRequest = requests[1]
         eq_(4, len(request.params))
         eq_("1", request.params["a0"])
         eq_("test_user", request.params["a1"])

packages/sqlalchemy-spanner/tests/mockserver_tests/test_quickstart.py

@@ -19,7 +19,6 @@
     CreateSessionRequest,
     ExecuteBatchDmlRequest,
     CommitRequest,
-    BeginTransactionRequest,
 )
 from sqlalchemy.orm import Session
 from sqlalchemy.testing import eq_, is_instance_of, is_not_none
@@ -115,12 +114,12 @@ def test_insert_data(self):
             session.commit()
 
             requests = self.spanner_service.requests
-            eq_(5, len(requests))
+            # Dialect now inlines BeginTransaction into the first statement.
+            eq_(4, len(requests))
             is_instance_of(requests[0], CreateSessionRequest)
-            is_instance_of(requests[1], BeginTransactionRequest)
+            is_instance_of(requests[1], ExecuteBatchDmlRequest)
             is_instance_of(requests[2], ExecuteBatchDmlRequest)
-            is_instance_of(requests[3], ExecuteBatchDmlRequest)
-            is_instance_of(requests[4], CommitRequest)
-            is_not_none(requests[2].transaction.id)
-            eq_(requests[2].transaction.id, requests[3].transaction.id)
-            eq_(requests[2].transaction.id, requests[4].transaction_id)
+            is_instance_of(requests[3], CommitRequest)
+            is_not_none(requests[1].transaction.begin)  # First request inlines begin
+            is_not_none(requests[2].transaction.id)  # Subsequent requests use ID
+            eq_(requests[2].transaction.id, requests[3].transaction_id)

packages/sqlalchemy-spanner/tests/mockserver_tests/test_tags.py

@@ -83,18 +83,18 @@ def test_transaction_tag(self):
 
         # Verify the requests that we got.
         requests = self.spanner_service.requests
-        eq_(6, len(requests))
+        # Dialect now inlines BeginTransaction into the first statement.
+        eq_(5, len(requests))
         is_instance_of(requests[0], CreateSessionRequest)
-        is_instance_of(requests[1], BeginTransactionRequest)
+        is_instance_of(requests[1], ExecuteSqlRequest)
         is_instance_of(requests[2], ExecuteSqlRequest)
         is_instance_of(requests[3], ExecuteSqlRequest)
-        is_instance_of(requests[4], ExecuteSqlRequest)
-        is_instance_of(requests[5], CommitRequest)
-        for request in requests[2:]:
+        is_instance_of(requests[4], CommitRequest)
+        for request in requests[1:]:
             eq_("my-transaction-tag", request.request_options.transaction_tag)
-        eq_("my-tag-1", requests[2].request_options.request_tag)
-        eq_("my-tag-2", requests[3].request_options.request_tag)
-        eq_("insert-singer", requests[4].request_options.request_tag)
+        eq_("my-tag-1", requests[1].request_options.request_tag)
+        eq_("my-tag-2", requests[2].request_options.request_tag)
+        eq_("insert-singer", requests[3].request_options.request_tag)
 
 
 def empty_singer_result_set():