googleapis · olavloite · Jun 3, 2025 · Jun 2, 2025 · Jun 2, 2025 · Jun 3, 2025
diff --git a/.github/workflows/mockserver-tests.yml b/.github/workflows/mockserver-tests.yml
@@ -0,0 +1,21 @@
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+name: Run Django Spanner mockserver tests
+jobs:
+  mockserver-tests:
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+      - name: Set up Python 3.12
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+      - name: Install nox
+        run: python -m pip install nox
+      - name: Run nox
+        run: nox -s mockserver
diff --git a/django_spanner/base.py b/django_spanner/base.py
@@ -120,9 +120,11 @@ def instance(self):
         :rtype: :class:`~google.cloud.spanner_v1.instance.Instance`
         :returns: A new instance owned by the existing Spanner Client.
         """
-        return spanner.Client(
-            project=os.environ["GOOGLE_CLOUD_PROJECT"]
-        ).instance(self.settings_dict["INSTANCE"])
+        if "client" in self.settings_dict["OPTIONS"]:
+            client = self.settings_dict["OPTIONS"]["client"]
+        else:
+            client = spanner.Client(project=os.environ["GOOGLE_CLOUD_PROJECT"])
+        return client.instance(self.settings_dict["INSTANCE"])
 
     @property
     def allow_transactions_in_auto_commit(self):

diff --git a/noxfile.py b/noxfile.py
@@ -24,6 +24,7 @@
     "setup.py",
 ]
 
+MOCKSERVER_TEST_PYTHON_VERSION = "3.12"
 DEFAULT_PYTHON_VERSION = "3.8"
 SYSTEM_TEST_PYTHON_VERSIONS = ["3.8"]
 UNIT_TEST_PYTHON_VERSIONS = ["3.8", "3.9", "3.10"]
@@ -69,6 +70,7 @@ def lint_setup_py(session):
 def default(session, django_version="3.2"):
     # Install all test dependencies, then install this package in-place.
     session.install(
+        "setuptools",
         "django~={}".format(django_version),
         "mock",
         "mock-import",
@@ -107,6 +109,32 @@ def unit(session):
     default(session, django_version="4.2")
 
 
+@nox.session(python=MOCKSERVER_TEST_PYTHON_VERSION)
+def mockserver(session):
+    # Install all test dependencies, then install this package in-place.
+    session.install(
+        "setuptools",
+        "django~=4.2",
+        "mock",
+        "mock-import",
+        "pytest",
+        "pytest-cov",
+        "coverage",
+        "sqlparse>=0.4.4",
+        "google-cloud-spanner>=3.55.0",
+        "opentelemetry-api==1.1.0",
+        "opentelemetry-sdk==1.1.0",
+        "opentelemetry-instrumentation==0.20b0",
+    )
+    session.install("-e", ".")
+    session.run(
+        "py.test",
+        "--quiet",
+        os.path.join("tests", "mockserver_tests"),
+        *session.posargs,
+    )
+
+
 def system_test(session, django_version="3.2"):
     """Run the system test suite."""
     constraints_path = str(

diff --git a/tests/mockserver_tests/__init__.py b/tests/mockserver_tests/__init__.py
diff --git a/tests/mockserver_tests/mock_database_admin.py b/tests/mockserver_tests/mock_database_admin.py
@@ -0,0 +1,38 @@
+# Copyright 2025 Google LLC All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from google.protobuf import empty_pb2
+import tests.mockserver_tests.spanner_database_admin_pb2_grpc as database_admin_grpc
+from google.longrunning import operations_pb2 as operations_pb2
+
+
+# An in-memory mock DatabaseAdmin server that can be used for testing.
+class DatabaseAdminServicer(database_admin_grpc.DatabaseAdminServicer):
+    def __init__(self):
+        self._requests = []
+
+    @property
+    def requests(self):
+        return self._requests
+
+    def clear_requests(self):
+        self._requests = []
+
+    def UpdateDatabaseDdl(self, request, context):
+        self._requests.append(request)
+        operation = operations_pb2.Operation()
+        operation.done = True
+        operation.name = "projects/test-project/operations/test-operation"
+        operation.response.Pack(empty_pb2.Empty())
+        return operation
diff --git a/tests/mockserver_tests/mock_server_test_base.py b/tests/mockserver_tests/mock_server_test_base.py
@@ -0,0 +1,223 @@
+# Copyright 2025 Google LLC All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import unittest
+
+from django.db import connection
+from google.cloud.spanner_dbapi.parsed_statement import AutocommitDmlMode
+import google.cloud.spanner_v1.types.type as spanner_type
+import google.cloud.spanner_v1.types.result_set as result_set
+from google.api_core.client_options import ClientOptions
+from google.auth.credentials import AnonymousCredentials
+from google.cloud.spanner_v1 import (
+    Client,
+    ResultSet,
+    PingingPool,
+    TypeCode,
+)
+from google.cloud.spanner_v1.database import Database
+from google.cloud.spanner_v1.instance import Instance
+import grpc
+
+# TODO: Replace this with the mock server in the Spanner client lib
+from tests.mockserver_tests.mock_spanner import (
+    SpannerServicer,
+    start_mock_server,
+)
+from tests.mockserver_tests.mock_database_admin import DatabaseAdminServicer
+
+
+def add_result(sql: str, result: 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",)])
+
+
+def add_single_result(
+    sql: str, column_name: str, type_code: spanner_type.TypeCode, row
+):
+    result = result_set.ResultSet(
+        dict(
+            metadata=result_set.ResultSetMetadata(
+                dict(
+                    row_type=spanner_type.StructType(
+                        dict(
+                            fields=[
+                                spanner_type.StructType.Field(
+                                    dict(
+                                        name=column_name,
+                                        type=spanner_type.Type(
+                                            dict(code=type_code)
+                                        ),
+                                    )
+                                )
+                            ]
+                        )
+                    )
+                )
+            ),
+        )
+    )
+    result.rows.extend(row)
+    MockServerTestBase.spanner_service.mock_spanner.add_result(sql, result)
+
+
+def add_singer_query_result(sql: str):
+    result = result_set.ResultSet(
+        dict(
+            metadata=result_set.ResultSetMetadata(
+                dict(
+                    row_type=spanner_type.StructType(
+                        dict(
+                            fields=[
+                                spanner_type.StructType.Field(
+                                    dict(
+                                        name="id",
+                                        type=spanner_type.Type(
+                                            dict(
+                                                code=spanner_type.TypeCode.INT64
+                                            )
+                                        ),
+                                    )
+                                ),
+                                spanner_type.StructType.Field(
+                                    dict(
+                                        name="first_name",
+                                        type=spanner_type.Type(
+                                            dict(
+                                                code=spanner_type.TypeCode.STRING
+                                            )
+                                        ),
+                                    )
+                                ),
+                                spanner_type.StructType.Field(
+                                    dict(
+                                        name="last_name",
+                                        type=spanner_type.Type(
+                                            dict(
+                                                code=spanner_type.TypeCode.STRING
+                                            )
+                                        ),
+                                    )
+                                ),
+                            ]
+                        )
+                    )
+                )
+            ),
+        )
+    )
+    result.rows.extend(
+        [
+            (
+                "1",
+                "Jane",
+                "Doe",
+            ),
+            (
+                "2",
+                "John",
+                "Doe",
+            ),
+        ]
+    )
+    add_result(sql, result)
+
+
+class MockServerTestBase(unittest.TestCase):
+    server: grpc.Server = None
+    spanner_service: SpannerServicer = None
+    database_admin_service: DatabaseAdminServicer = None
+    port: int = None
+    _client = None
+    _instance = None
+    _database = None
+    _pool = None
+
+    @classmethod
+    def setup_class(cls):
+        os.environ["GOOGLE_CLOUD_PROJECT"] = "mockserver-project"
+        (
+            MockServerTestBase.server,
+            MockServerTestBase.spanner_service,
+            MockServerTestBase.database_admin_service,
+            MockServerTestBase.port,
+        ) = start_mock_server()
+
+    @classmethod
+    def teardown_class(cls):
+        if MockServerTestBase.server is not None:
+            MockServerTestBase.server.stop(grace=None)
+            MockServerTestBase.server = None
+
+    def setup_method(self, test_method):
+        connection.settings_dict["OPTIONS"]["client"] = self.client
+        connection.settings_dict["OPTIONS"]["pool"] = self.pool
+
+    def teardown_method(self, test_method):
+        connection.close()
+        MockServerTestBase.spanner_service.clear_requests()
+        MockServerTestBase.database_admin_service.clear_requests()
+        self._client = None
+        self._instance = None
+        self._database = None
+        self._pool = None
+
+    @property
+    def client(self) -> Client:
+        if self._client is None:
+            self._client = Client(
+                project=os.environ["GOOGLE_CLOUD_PROJECT"],
+                credentials=AnonymousCredentials(),
+                client_options=ClientOptions(
+                    api_endpoint="localhost:" + str(MockServerTestBase.port),
+                ),
+            )
+        return self._client
+
+    @property
+    def pool(self):
+        if self._pool is None:
+            self._pool = PingingPool(size=10)
+        return self._pool
+
+    @property
+    def instance(self) -> Instance:
+        if self._instance is None:
+            self._instance = self.client.instance("test-instance")
+        return self._instance
+
+    @property
+    def database(self) -> Database:
+        if self._database is None:
+            self._database = self.instance.database(
+                "test-database", pool=PingingPool(size=10)
+            )
+        return self._database