Skip to content

Commit 96da8e1

Browse files
olavloitegcf-owl-bot[bot]
olavloite
and
gcf-owl-bot[bot]
authored
test: create base class for mockserver tests (#1255)
* test: create base class for mockserver tests Move the boiler-plate code for mockserver tests to a separate class, so this can easily be re-used for other tests. * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md --------- Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
1 parent 758bf48 commit 96da8e1

File tree

2 files changed

+147
-113
lines changed

2 files changed

+147
-113
lines changed

‎tests/mockserver_tests/mock_server_test_base.py‎

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
1+
# Copyright 2024 Google LLC All rights reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import unittest
16+
17+
from google.cloud.spanner_dbapi.parsed_statement import AutocommitDmlMode
18+
from google.cloud.spanner_v1.testing.mock_database_admin import DatabaseAdminServicer
19+
from google.cloud.spanner_v1.testing.mock_spanner import (
20+
start_mock_server,
21+
SpannerServicer,
22+
)
23+
import google.cloud.spanner_v1.types.type as spanner_type
24+
import google.cloud.spanner_v1.types.result_set as result_set
25+
from google.api_core.client_options import ClientOptions
26+
from google.auth.credentials import AnonymousCredentials
27+
from google.cloud.spanner_v1 import Client, TypeCode, FixedSizePool
28+
from google.cloud.spanner_v1.database import Database
29+
from google.cloud.spanner_v1.instance import Instance
30+
import grpc
31+
32+
33+
def add_result(sql: str, result: result_set.ResultSet):
34+
MockServerTestBase.spanner_service.mock_spanner.add_result(sql, result)
35+
36+
37+
def add_update_count(
38+
sql: str, count: int, dml_mode: AutocommitDmlMode = AutocommitDmlMode.TRANSACTIONAL
39+
):
40+
if dml_mode == AutocommitDmlMode.PARTITIONED_NON_ATOMIC:
41+
stats = dict(row_count_lower_bound=count)
42+
else:
43+
stats = dict(row_count_exact=count)
44+
result = result_set.ResultSet(dict(stats=result_set.ResultSetStats(stats)))
45+
add_result(sql, result)
46+
47+
48+
def add_select1_result():
49+
add_single_result("select 1", "c", TypeCode.INT64, [("1",)])
50+
51+
52+
def add_single_result(
53+
sql: str, column_name: str, type_code: spanner_type.TypeCode, row
54+
):
55+
result = result_set.ResultSet(
56+
dict(
57+
metadata=result_set.ResultSetMetadata(
58+
dict(
59+
row_type=spanner_type.StructType(
60+
dict(
61+
fields=[
62+
spanner_type.StructType.Field(
63+
dict(
64+
name=column_name,
65+
type=spanner_type.Type(dict(code=type_code)),
66+
)
67+
)
68+
]
69+
)
70+
)
71+
)
72+
),
73+
)
74+
)
75+
result.rows.extend(row)
76+
MockServerTestBase.spanner_service.mock_spanner.add_result(sql, result)
77+
78+
79+
class MockServerTestBase(unittest.TestCase):
80+
server: grpc.Server = None
81+
spanner_service: SpannerServicer = None
82+
database_admin_service: DatabaseAdminServicer = None
83+
port: int = None
84+
85+
def __init__(self, *args, **kwargs):
86+
super(MockServerTestBase, self).__init__(*args, **kwargs)
87+
self._client = None
88+
self._instance = None
89+
self._database = None
90+
91+
@classmethod
92+
def setup_class(cls):
93+
(
94+
MockServerTestBase.server,
95+
MockServerTestBase.spanner_service,
96+
MockServerTestBase.database_admin_service,
97+
MockServerTestBase.port,
98+
) = start_mock_server()
99+
100+
@classmethod
101+
def teardown_class(cls):
102+
if MockServerTestBase.server is not None:
103+
MockServerTestBase.server.stop(grace=None)
104+
MockServerTestBase.server = None
105+
106+
def setup_method(self, *args, **kwargs):
107+
self._client = None
108+
self._instance = None
109+
self._database = None
110+
111+
def teardown_method(self, *args, **kwargs):
112+
MockServerTestBase.spanner_service.clear_requests()
113+
MockServerTestBase.database_admin_service.clear_requests()
114+
115+
@property
116+
def client(self) -> Client:
117+
if self._client is None:
118+
self._client = Client(
119+
project="p",
120+
credentials=AnonymousCredentials(),
121+
client_options=ClientOptions(
122+
api_endpoint="localhost:" + str(MockServerTestBase.port),
123+
),
124+
)
125+
return self._client
126+
127+
@property
128+
def instance(self) -> Instance:
129+
if self._instance is None:
130+
self._instance = self.client.instance("test-instance")
131+
return self._instance
132+
133+
@property
134+
def database(self) -> Database:
135+
if self._database is None:
136+
self._database = self.instance.database(
137+
"test-database", pool=FixedSizePool(size=10)
138+
)
139+
return self._database

‎tests/mockserver_tests/test_basics.py‎

Lines changed: 8 additions & 113 deletions
Original file line numberDiff line numberDiff line change
@@ -12,131 +12,26 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15-
import unittest
16-
1715
from google.cloud.spanner_admin_database_v1.types import spanner_database_admin
1816
from google.cloud.spanner_dbapi import Connection
1917
from google.cloud.spanner_dbapi.parsed_statement import AutocommitDmlMode
20-
from google.cloud.spanner_v1.testing.mock_database_admin import DatabaseAdminServicer
21-
from google.cloud.spanner_v1.testing.mock_spanner import (
22-
start_mock_server,
23-
SpannerServicer,
24-
)
25-
import google.cloud.spanner_v1.types.type as spanner_type
26-
import google.cloud.spanner_v1.types.result_set as result_set
27-
from google.api_core.client_options import ClientOptions
28-
from google.auth.credentials import AnonymousCredentials
2918
from google.cloud.spanner_v1 import (
30-
Client,
31-
FixedSizePool,
3219
BatchCreateSessionsRequest,
3320
ExecuteSqlRequest,
3421
BeginTransactionRequest,
3522
TransactionOptions,
3623
)
37-
from google.cloud.spanner_v1.database import Database
38-
from google.cloud.spanner_v1.instance import Instance
39-
import grpc
40-
41-
42-
class TestBasics(unittest.TestCase):
43-
server: grpc.Server = None
44-
spanner_service: SpannerServicer = None
45-
database_admin_service: DatabaseAdminServicer = None
46-
port: int = None
47-
48-
def __init__(self, *args, **kwargs):
49-
super(TestBasics, self).__init__(*args, **kwargs)
50-
self._client = None
51-
self._instance = None
52-
self._database = None
5324

54-
@classmethod
55-
def setUpClass(cls):
56-
(
57-
TestBasics.server,
58-
TestBasics.spanner_service,
59-
TestBasics.database_admin_service,
60-
TestBasics.port,
61-
) = start_mock_server()
62-
63-
@classmethod
64-
def tearDownClass(cls):
65-
if TestBasics.server is not None:
66-
TestBasics.server.stop(grace=None)
67-
TestBasics.server = None
68-
69-
def teardown_method(self, *args, **kwargs):
70-
TestBasics.spanner_service.clear_requests()
71-
TestBasics.database_admin_service.clear_requests()
72-
73-
def _add_select1_result(self):
74-
result = result_set.ResultSet(
75-
dict(
76-
metadata=result_set.ResultSetMetadata(
77-
dict(
78-
row_type=spanner_type.StructType(
79-
dict(
80-
fields=[
81-
spanner_type.StructType.Field(
82-
dict(
83-
name="c",
84-
type=spanner_type.Type(
85-
dict(code=spanner_type.TypeCode.INT64)
86-
),
87-
)
88-
)
89-
]
90-
)
91-
)
92-
)
93-
),
94-
)
95-
)
96-
result.rows.extend(["1"])
97-
TestBasics.spanner_service.mock_spanner.add_result("select 1", result)
98-
99-
def add_update_count(
100-
self,
101-
sql: str,
102-
count: int,
103-
dml_mode: AutocommitDmlMode = AutocommitDmlMode.TRANSACTIONAL,
104-
):
105-
if dml_mode == AutocommitDmlMode.PARTITIONED_NON_ATOMIC:
106-
stats = dict(row_count_lower_bound=count)
107-
else:
108-
stats = dict(row_count_exact=count)
109-
result = result_set.ResultSet(dict(stats=result_set.ResultSetStats(stats)))
110-
TestBasics.spanner_service.mock_spanner.add_result(sql, result)
111-
112-
@property
113-
def client(self) -> Client:
114-
if self._client is None:
115-
self._client = Client(
116-
project="test-project",
117-
credentials=AnonymousCredentials(),
118-
client_options=ClientOptions(
119-
api_endpoint="localhost:" + str(TestBasics.port),
120-
),
121-
)
122-
return self._client
123-
124-
@property
125-
def instance(self) -> Instance:
126-
if self._instance is None:
127-
self._instance = self.client.instance("test-instance")
128-
return self._instance
25+
from tests.mockserver_tests.mock_server_test_base import (
26+
MockServerTestBase,
27+
add_select1_result,
28+
add_update_count,
29+
)
12930

130-
@property
131-
def database(self) -> Database:
132-
if self._database is None:
133-
self._database = self.instance.database(
134-
"test-database", pool=FixedSizePool(size=10)
135-
)
136-
return self._database
13731

32+
class TestBasics(MockServerTestBase):
13833
def test_select1(self):
139-
self._add_select1_result()
34+
add_select1_result()
14035
with self.database.snapshot() as snapshot:
14136
results = snapshot.execute_sql("select 1")
14237
result_list = []
@@ -171,7 +66,7 @@ def test_create_table(self):
17166
# been re-factored to use a base class for the boiler plate code.
17267
def test_dbapi_partitioned_dml(self):
17368
sql = "UPDATE singers SET foo='bar' WHERE active = true"
174-
self.add_update_count(sql, 100, AutocommitDmlMode.PARTITIONED_NON_ATOMIC)
69+
add_update_count(sql, 100, AutocommitDmlMode.PARTITIONED_NON_ATOMIC)
17570
connection = Connection(self.instance, self.database)
17671
connection.autocommit = True
17772
connection.set_autocommit_dml_mode(AutocommitDmlMode.PARTITIONED_NON_ATOMIC)

0 commit comments

Comments
 (0)