Skip to content

Commit 53a909a

Browse files
authored
test: adds test file (#2272)
* add copperpenny client. tests to follow in a separate PR * removes old copperpenny client * blackens and lints * removes obsolete test file. new test file to be added in future PR * test: adds test file
1 parent 0fc6877 commit 53a909a

File tree

1 file changed

+332
-0
lines changed

1 file changed

+332
-0
lines changed
Lines changed: 332 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,332 @@
1+
# -*- coding: utf-8 -*-
2+
# Copyright 2025 Google LLC
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
#
16+
17+
import pytest
18+
from typing import (
19+
Optional,
20+
Sequence,
21+
Tuple,
22+
Union,
23+
)
24+
from unittest import mock
25+
26+
from google.api_core import client_options as client_options_lib
27+
from google.api_core import gapic_v1
28+
from google.api_core import retry as retries
29+
from google.auth import credentials as auth_credentials
30+
31+
# --- IMPORT SERVICECLIENT MODULES ---
32+
from google.cloud.bigquery_v2.services import (
33+
centralized_service,
34+
dataset_service,
35+
job_service,
36+
model_service,
37+
)
38+
39+
# --- IMPORT TYPES MODULES (to access *Requests classes) ---
40+
from google.cloud.bigquery_v2.types import (
41+
dataset,
42+
job,
43+
model,
44+
)
45+
46+
# --- TYPE ALIASES ---
47+
try:
48+
OptionalRetry = Union[retries.Retry, gapic_v1.method._MethodDefault, None]
49+
except AttributeError: # pragma: NO COVER
50+
OptionalRetry = Union[retries.Retry, object, None] # type: ignore
51+
52+
AnyRequest = Union[
53+
dataset.GetDatasetRequest,
54+
model.GetModelRequest,
55+
model.DeleteModelRequest,
56+
model.PatchModelRequest,
57+
job.ListJobsRequest,
58+
model.ListModelsRequest,
59+
]
60+
61+
# --- CONSTANTS ---
62+
PROJECT_ID = "test-project"
63+
DATASET_ID = "test_dataset"
64+
JOB_ID = "test_job"
65+
MODEL_ID = "test_model"
66+
DEFAULT_ETAG = "test_etag"
67+
68+
DEFAULT_RETRY: OptionalRetry = gapic_v1.method.DEFAULT
69+
DEFAULT_TIMEOUT: Union[float, object] = gapic_v1.method.DEFAULT
70+
DEFAULT_METADATA: Sequence[Tuple[str, Union[str, bytes]]] = ()
71+
72+
# --- HELPERS ---
73+
def assert_client_called_once_with(
74+
mock_method: mock.Mock,
75+
request: AnyRequest,
76+
retry: OptionalRetry = DEFAULT_RETRY,
77+
timeout: Union[float, object] = DEFAULT_TIMEOUT,
78+
metadata: Sequence[Tuple[str, Union[str, bytes]]] = DEFAULT_METADATA,
79+
):
80+
"""Helper to assert a client method was called with default args."""
81+
mock_method.assert_called_once_with(
82+
request=request,
83+
retry=retry,
84+
timeout=timeout,
85+
metadata=metadata,
86+
)
87+
88+
89+
# --- FIXTURES ---
90+
@pytest.fixture
91+
def mock_dataset_service_client():
92+
"""Mocks the DatasetServiceClient."""
93+
with mock.patch(
94+
"google.cloud.bigquery_v2.services.dataset_service.DatasetServiceClient",
95+
autospec=True,
96+
) as mock_client:
97+
yield mock_client
98+
99+
100+
@pytest.fixture
101+
def mock_job_service_client():
102+
"""Mocks the JobServiceClient."""
103+
with mock.patch(
104+
"google.cloud.bigquery_v2.services.job_service.JobServiceClient",
105+
autospec=True,
106+
) as mock_client:
107+
yield mock_client
108+
109+
110+
@pytest.fixture
111+
def mock_model_service_client():
112+
"""Mocks the ModelServiceClient."""
113+
with mock.patch(
114+
"google.cloud.bigquery_v2.services.model_service.ModelServiceClient",
115+
autospec=True,
116+
) as mock_client:
117+
yield mock_client
118+
119+
120+
# TODO: figure out a solution for this... is there an easier way to feed in clients?
121+
# TODO: is there an easier way to make mock_x_service_clients?
122+
@pytest.fixture
123+
def bq_client(
124+
mock_dataset_service_client, mock_job_service_client, mock_model_service_client
125+
):
126+
"""Provides a BigQueryClient with mocked underlying services."""
127+
client = centralized_service.BigQueryClient()
128+
client.dataset_service_client = mock_dataset_service_client
129+
client.job_service_client = mock_job_service_client
130+
client.model_service_client = mock_model_service_client
131+
...
132+
return client
133+
134+
135+
# --- TEST CLASSES ---
136+
137+
from google.api_core import client_options as client_options_lib
138+
139+
# from google.api_core.client_options import ClientOptions
140+
from google.auth import credentials as auth_credentials
141+
142+
# from google.auth.credentials import Credentials
143+
144+
145+
class TestCentralizedClientInitialization:
146+
@pytest.mark.parametrize(
147+
"credentials, client_options",
148+
[
149+
(None, None),
150+
(mock.MagicMock(spec=auth_credentials.Credentials), None),
151+
(
152+
None,
153+
client_options_lib.ClientOptions(api_endpoint="test.googleapis.com"),
154+
),
155+
(
156+
mock.MagicMock(spec=auth_credentials.Credentials),
157+
client_options_lib.ClientOptions(api_endpoint="test.googleapis.com"),
158+
),
159+
],
160+
)
161+
def test_client_initialization_arguments(
162+
self,
163+
credentials,
164+
client_options,
165+
mock_dataset_service_client,
166+
mock_job_service_client,
167+
mock_model_service_client,
168+
):
169+
# Act
170+
client = centralized_service.BigQueryClient(
171+
credentials=credentials, client_options=client_options
172+
)
173+
174+
# Assert
175+
# The BigQueryClient should have been initialized. Accessing the
176+
# service client properties should instantiate them with the correct arguments.
177+
178+
# Access the property to trigger instantiation
179+
_ = client.dataset_service_client
180+
mock_dataset_service_client.assert_called_once_with(
181+
credentials=credentials, client_options=client_options
182+
)
183+
184+
_ = client.job_service_client
185+
mock_job_service_client.assert_called_once_with(
186+
credentials=credentials, client_options=client_options
187+
)
188+
189+
_ = client.model_service_client
190+
mock_model_service_client.assert_called_once_with(
191+
credentials=credentials, client_options=client_options
192+
)
193+
194+
195+
class TestCentralizedClientDatasetService:
196+
def test_get_dataset(self, bq_client, mock_dataset_service_client):
197+
# Arrange
198+
expected_dataset = dataset.Dataset(
199+
kind="bigquery#dataset", id=f"{PROJECT_ID}:{DATASET_ID}"
200+
)
201+
mock_dataset_service_client.get_dataset.return_value = expected_dataset
202+
get_dataset_request = dataset.GetDatasetRequest(
203+
project_id=PROJECT_ID, dataset_id=DATASET_ID
204+
)
205+
206+
# Act
207+
dataset_response = bq_client.get_dataset(request=get_dataset_request)
208+
209+
# Assert
210+
assert dataset_response == expected_dataset
211+
assert_client_called_once_with(
212+
mock_dataset_service_client.get_dataset, get_dataset_request
213+
)
214+
215+
216+
class TestCentralizedClientJobService:
217+
def test_list_jobs(self, bq_client, mock_job_service_client):
218+
# Arrange
219+
expected_jobs = [job.Job(kind="bigquery#job", id=f"{PROJECT_ID}:{JOB_ID}")]
220+
mock_job_service_client.list_jobs.return_value = expected_jobs
221+
list_jobs_request = job.ListJobsRequest(project_id=PROJECT_ID)
222+
223+
# Act
224+
jobs_response = bq_client.list_jobs(request=list_jobs_request)
225+
226+
# Assert
227+
assert jobs_response == expected_jobs
228+
assert_client_called_once_with(
229+
mock_job_service_client.list_jobs, list_jobs_request
230+
)
231+
232+
233+
class TestCentralizedClientModelService:
234+
def test_get_model(self, bq_client, mock_model_service_client):
235+
# Arrange
236+
expected_model = model.Model(
237+
etag=DEFAULT_ETAG,
238+
model_reference={
239+
"project_id": PROJECT_ID,
240+
"dataset_id": DATASET_ID,
241+
"model_id": MODEL_ID,
242+
},
243+
)
244+
mock_model_service_client.get_model.return_value = expected_model
245+
get_model_request = model.GetModelRequest(
246+
project_id=PROJECT_ID, dataset_id=DATASET_ID, model_id=MODEL_ID
247+
)
248+
249+
# Act
250+
model_response = bq_client.get_model(request=get_model_request)
251+
252+
# Assert
253+
assert model_response == expected_model
254+
assert_client_called_once_with(
255+
mock_model_service_client.get_model, get_model_request
256+
)
257+
258+
def test_delete_model(self, bq_client, mock_model_service_client):
259+
# Arrange
260+
# The underlying service call returns nothing on success.
261+
mock_model_service_client.delete_model.return_value = None
262+
delete_model_request = model.DeleteModelRequest(
263+
project_id=PROJECT_ID, dataset_id=DATASET_ID, model_id=MODEL_ID
264+
)
265+
266+
# Act
267+
# The wrapper method should also return nothing.
268+
result = bq_client.delete_model(request=delete_model_request)
269+
270+
# Assert
271+
# 1. Assert the return value is None. This fails if the method doesn't exist.
272+
assert result is None
273+
# 2. Assert the underlying service was called correctly.
274+
assert_client_called_once_with(
275+
mock_model_service_client.delete_model,
276+
delete_model_request,
277+
)
278+
279+
def test_patch_model(self, bq_client, mock_model_service_client):
280+
# Arrange
281+
expected_model = model.Model(
282+
etag="new_etag",
283+
model_reference={
284+
"project_id": PROJECT_ID,
285+
"dataset_id": DATASET_ID,
286+
"model_id": MODEL_ID,
287+
},
288+
description="A newly patched description.",
289+
)
290+
mock_model_service_client.patch_model.return_value = expected_model
291+
292+
model_patch = model.Model(description="A newly patched description.")
293+
patch_model_request = model.PatchModelRequest(
294+
project_id=PROJECT_ID,
295+
dataset_id=DATASET_ID,
296+
model_id=MODEL_ID,
297+
model=model_patch,
298+
)
299+
300+
# Act
301+
patched_model = bq_client.patch_model(request=patch_model_request)
302+
303+
# Assert
304+
assert patched_model == expected_model
305+
assert_client_called_once_with(
306+
mock_model_service_client.patch_model, patch_model_request
307+
)
308+
309+
def test_list_models(self, bq_client, mock_model_service_client):
310+
# Arrange
311+
expected_models = [
312+
model.Model(
313+
etag=DEFAULT_ETAG,
314+
model_reference={
315+
"project_id": PROJECT_ID,
316+
"dataset_id": DATASET_ID,
317+
"model_id": MODEL_ID,
318+
},
319+
)
320+
]
321+
mock_model_service_client.list_models.return_value = expected_models
322+
list_models_request = model.ListModelsRequest(
323+
project_id=PROJECT_ID, dataset_id=DATASET_ID
324+
)
325+
# Act
326+
models_response = bq_client.list_models(request=list_models_request)
327+
328+
# Assert
329+
assert models_response == expected_models
330+
assert_client_called_once_with(
331+
mock_model_service_client.list_models, list_models_request
332+
)

0 commit comments

Comments
 (0)