Merge branch 'main' into make-base-mock-server-test-class

olavloite · Dec 5, 2024 · Dec 5, 2024 · Dec 5, 2024 · Dec 5, 2024 · Dec 5, 2024
commit f4ca07306d01de639c1dfb223dddd653551cd422
diff --git a/noxfile.py b/noxfile.py
@@ -32,7 +32,7 @@
 ISORT_VERSION = "isort==5.11.0"
 LINT_PATHS = ["docs", "google", "tests", "noxfile.py", "setup.py"]
 
-DEFAULT_PYTHON_VERSION = "3.8"
+DEFAULT_PYTHON_VERSION = "3.12"
 
 DEFAULT_MOCK_SERVER_TESTS_PYTHON_VERSION = "3.12"
 UNIT_TEST_PYTHON_VERSIONS: List[str] = [

@@ -14,6 +14,7 @@
 
 import unittest
 
+from google.cloud.spanner_dbapi.parsed_statement import AutocommitDmlMode
 from google.cloud.spanner_v1.testing.mock_database_admin import DatabaseAdminServicer
 from google.cloud.spanner_v1.testing.mock_spanner import (
     start_mock_server,
@@ -33,6 +34,17 @@ def add_result(sql: str, result: result_set.ResultSet):
     MockServerTestBase.spanner_service.mock_spanner.add_result(sql, result)
 
 
+def add_update_count(
+    sql: str, count: int, dml_mode: AutocommitDmlMode = AutocommitDmlMode.TRANSACTIONAL
+):
+    if dml_mode == AutocommitDmlMode.PARTITIONED_NON_ATOMIC:
+        stats = dict(row_count_lower_bound=count)
+    else:
+        stats = dict(row_count_exact=count)
+    result = result_set.ResultSet(dict(stats=result_set.ResultSetStats(stats)))
+    add_result(sql, result)
+
+
 def add_select1_result():
     add_single_result("select 1", "c", TypeCode.INT64, [("1",)])
 

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 from google.cloud.spanner_admin_database_v1.types import spanner_database_admin
+from google.cloud.spanner_dbapi import Connection
+from google.cloud.spanner_dbapi.parsed_statement import AutocommitDmlMode
 from google.cloud.spanner_v1 import (
     BatchCreateSessionsRequest,
     ExecuteSqlRequest,
@@ -23,6 +25,7 @@
 from tests.mockserver_tests.mock_server_test_base import (
     MockServerTestBase,
     add_select1_result,
+    add_update_count,
 )
 
 
@@ -63,7 +66,7 @@ def test_create_table(self):
     #       been re-factored to use a base class for the boiler plate code.
     def test_dbapi_partitioned_dml(self):
         sql = "UPDATE singers SET foo='bar' WHERE active = true"
-        self.add_update_count(sql, 100, AutocommitDmlMode.PARTITIONED_NON_ATOMIC)
+        add_update_count(sql, 100, AutocommitDmlMode.PARTITIONED_NON_ATOMIC)
         connection = Connection(self.instance, self.database)
         connection.autocommit = True
         connection.set_autocommit_dml_mode(AutocommitDmlMode.PARTITIONED_NON_ATOMIC)