add mockspanner tests

Author: sakthivelmanii

google/cloud/spanner_v1/testing/mock_spanner.py

@@ -35,11 +35,17 @@
 class MockSpanner:
     def __init__(self):
         self.results = {}
+        self.execute_streaming_sql_results = {}
         self.errors = {}
 
     def add_result(self, sql: str, result: result_set.ResultSet):
         self.results[sql.lower().strip()] = result
 
+    def add_execute_streaming_sql_results(self, sql: str,
+        partial_result_sets: list[result_set.PartialResultSet]):
+        self.execute_streaming_sql_results[
+            sql.lower().strip()] = partial_result_sets
+
     def get_result(self, sql: str) -> result_set.ResultSet:
         result = self.results.get(sql.lower().strip())
         if result is None:
@@ -55,9 +61,20 @@ def pop_error(self, context):
         if error:
             context.abort_with_status(error)
 
+    def get_execute_streaming_sql_results(self, sql: str,
+        started_transaction: transaction.Transaction) -> list[
+        result_set.PartialResultSet]:
+        if self.execute_streaming_sql_results[sql.lower().strip()]:
+            partials = self.execute_streaming_sql_results[sql.lower().strip()]
+        else:
+            partials = self.get_result_as_partial_result_sets(sql)
+        if started_transaction:
+            partials[0].metadata.transaction = started_transaction
+        return partials
+
     def get_result_as_partial_result_sets(
-        self, sql: str, started_transaction: transaction.Transaction
-    ) -> [result_set.PartialResultSet]:
+        self, sql: str
+    ) -> list[result_set.PartialResultSet]:
         result: result_set.ResultSet = self.get_result(sql)
         partials = []
         first = True
@@ -70,11 +87,10 @@ def get_result_as_partial_result_sets(
                 partial = result_set.PartialResultSet()
                 if first:
                     partial.metadata = ResultSetMetadata(result.metadata)
+                first = False
                 partial.values.extend(row)
                 partials.append(partial)
         partials[len(partials) - 1].stats = result.stats
-        if started_transaction:
-            partials[0].metadata.transaction = started_transaction
         return partials
 
 
@@ -149,7 +165,7 @@ def ExecuteStreamingSql(self, request, context):
         self._requests.append(request)
         self.mock_spanner.pop_error(context)
         started_transaction = self.__maybe_create_transaction(request)
-        partials = self.mock_spanner.get_result_as_partial_result_sets(
+        partials = self.mock_spanner.get_execute_streaming_sql_results(
             request.sql, started_transaction
         )
         for result in partials:

tests/mockserver_tests/mock_server_test_base.py

@@ -14,27 +14,36 @@
 
 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,
-    SpannerServicer,
-)
-import google.cloud.spanner_v1.types.type as spanner_type
-import google.cloud.spanner_v1.types.result_set as result_set
+import grpc
 from google.api_core.client_options import ClientOptions
 from google.auth.credentials import AnonymousCredentials
-from google.cloud.spanner_v1 import Client, TypeCode, FixedSizePool
-from google.cloud.spanner_v1.database import Database
-from google.cloud.spanner_v1.instance import Instance
-import grpc
+from google.cloud.spanner_v1 import Type
+
+from google.cloud.spanner_v1 import StructType
+from google.cloud.spanner_v1._helpers import _make_value_pb
+
+from google.cloud.spanner_v1 import PartialResultSet
+from google.protobuf.duration_pb2 import Duration
 from google.rpc import code_pb2
 from google.rpc import status_pb2
 from google.rpc.error_details_pb2 import RetryInfo
-from google.protobuf.duration_pb2 import Duration
 from grpc_status._common import code_to_grpc_status_code
 from grpc_status.rpc_status import _Status
 
+import google.cloud.spanner_v1.types.result_set as result_set
+import google.cloud.spanner_v1.types.type as spanner_type
+from google.cloud.spanner_dbapi.parsed_statement import AutocommitDmlMode
+from google.cloud.spanner_v1 import Client
+from google.cloud.spanner_v1 import FixedSizePool
+from google.cloud.spanner_v1 import ResultSetMetadata
+from google.cloud.spanner_v1 import TypeCode
+from google.cloud.spanner_v1.database import Database
+from google.cloud.spanner_v1.instance import Instance
+from google.cloud.spanner_v1.testing.mock_database_admin import \
+    DatabaseAdminServicer
+from google.cloud.spanner_v1.testing.mock_spanner import SpannerServicer
+from google.cloud.spanner_v1.testing.mock_spanner import start_mock_server
+
 
 # Creates an aborted status with the smallest possible retry delay.
 def aborted_status() -> _Status:
@@ -57,6 +66,24 @@ def aborted_status() -> _Status:
     return status
 
 
+def _make_partial_result_sets(fields: list[tuple[str, TypeCode]],
+    results: list[dict]) -> list[result_set.PartialResultSet]:
+    partial_result_sets = []
+    for result in results:
+        partial_result_set = PartialResultSet()
+        if len(partial_result_sets) == 0:
+            # setting the metadata
+            metadata = ResultSetMetadata(row_type=StructType(fields=[]))
+            for field in fields:
+                metadata.row_type.fields.append(
+                    StructType.Field(name=field[0], type_=Type(code=field[1])))
+            partial_result_set.metadata = metadata
+        for value in result["values"]:
+            partial_result_set.values.append(_make_value_pb(value))
+        partial_result_set.last = result.get('last') or False
+        partial_result_sets.append(partial_result_set)
+    return partial_result_sets
+
 # Creates an UNAVAILABLE status with the smallest possible retry delay.
 def unavailable_status() -> _Status:
     error = status_pb2.Status(
@@ -101,6 +128,11 @@ def add_select1_result():
     add_single_result("select 1", "c", TypeCode.INT64, [("1",)])
 
 
+def add_execute_streaming_sql_results(sql: str,
+    partial_result_sets: list[result_set.PartialResultSet]):
+    MockServerTestBase.spanner_service.mock_spanner.add_execute_streaming_sql_results(
+        sql, partial_result_sets)
+
 def add_single_result(
     sql: str, column_name: str, type_code: spanner_type.TypeCode, row
 ):

tests/mockserver_tests/test_basics.py

@@ -15,28 +15,32 @@
 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,
-    BeginTransactionRequest,
-    TransactionOptions,
-    ExecuteBatchDmlRequest,
-    TypeCode,
-)
-from google.cloud.spanner_v1.transaction import Transaction
+from google.cloud.spanner_v1 import BatchCreateSessionsRequest
+from google.cloud.spanner_v1 import BeginTransactionRequest
+from google.cloud.spanner_v1 import ExecuteBatchDmlRequest
+from google.cloud.spanner_v1 import ExecuteSqlRequest
+from google.cloud.spanner_v1 import TransactionOptions
+from google.cloud.spanner_v1 import TypeCode
 from google.cloud.spanner_v1.testing.mock_spanner import SpannerServicer
-
-from tests.mockserver_tests.mock_server_test_base import (
-    MockServerTestBase,
-    add_select1_result,
-    add_update_count,
-    add_error,
-    unavailable_status,
-    add_single_result,
-)
+from google.cloud.spanner_v1.transaction import Transaction
+from tests.mockserver_tests.mock_server_test_base import MockServerTestBase
+from tests.mockserver_tests.mock_server_test_base import \
+    _make_partial_result_sets
+from tests.mockserver_tests.mock_server_test_base import add_error
+from tests.mockserver_tests.mock_server_test_base import \
+    add_execute_streaming_sql_results
+from tests.mockserver_tests.mock_server_test_base import add_select1_result
+from tests.mockserver_tests.mock_server_test_base import add_single_result
+from tests.mockserver_tests.mock_server_test_base import add_update_count
+from tests.mockserver_tests.mock_server_test_base import unavailable_status
 
 
 class TestBasics(MockServerTestBase):
+
+    def setUp(self):
+        super().setUp()
+        super().setup_class()
+
     def test_select1(self):
         add_select1_result()
         with self.database.snapshot() as snapshot:
@@ -176,6 +180,27 @@ def test_last_statement_query(self):
         self.assertEqual(1, len(requests), msg=requests)
         self.assertTrue(requests[0].last_statement, requests[0])
 
+    def test_execute_streaming_sql_last_field(self):
+        partial_result_sets = _make_partial_result_sets(
+            [("ID", TypeCode.INT64), ("NAME", TypeCode.STRING)],
+            [{"values": ["1", "ABC", "2", "DEF"]},
+             {"values": ["3", "GHI"], "last": True}])
+
+        sql = "select * from my_table"
+        add_execute_streaming_sql_results(sql, partial_result_sets)
+        count = 1
+        with self.database.snapshot() as snapshot:
+            results = snapshot.execute_sql(sql)
+            result_list = []
+            for row in results:
+                result_list.append(row)
+                self.assertEqual(count, row[0])
+                count += 1
+            self.assertEqual(3, len(result_list))
+        requests = self.spanner_service.requests
+        self.assertEqual(2, len(requests), msg=requests)
+        self.assertTrue(isinstance(requests[0], BatchCreateSessionsRequest))
+        self.assertTrue(isinstance(requests[1], ExecuteSqlRequest))
 
 def _execute_query(transaction: Transaction, sql: str):
     rows = transaction.execute_sql(sql, last_statement=True)