This repository was archived by the owner on Apr 1, 2026. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 65
Expand file tree
/
Copy pathtest_client.py
More file actions
3153 lines (2942 loc) · 140 KB
/
test_client.py
File metadata and controls
3153 lines (2942 loc) · 140 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
# Copyright 2024 Google LLC
#
# 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.
# This file is automatically generated by CrossSync. Do not edit manually.
from __future__ import annotations
import grpc
import asyncio
import re
import pytest
import mock
from google.cloud.bigtable.data import mutations
from google.auth.credentials import AnonymousCredentials
from google.cloud.bigtable_v2.types import ReadRowsResponse
from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery
from google.api_core import exceptions as core_exceptions
from google.api_core import client_options
from google.cloud.bigtable.data.exceptions import InvalidChunk
from google.cloud.bigtable.data.exceptions import _MutateRowsIncomplete
from google.cloud.bigtable.data.mutations import DeleteAllFromRow
from google.cloud.bigtable.data import TABLE_DEFAULT
from google.cloud.bigtable.data.read_modify_write_rules import IncrementRule
from google.cloud.bigtable.data.read_modify_write_rules import AppendValueRule
from google.cloud.bigtable_v2.types.bigtable import ExecuteQueryResponse
from google.cloud.bigtable.data._cross_sync import CrossSync
from tests.unit.data.execute_query.sql_helpers import (
chunked_responses,
column,
int64_type,
int_val,
metadata,
null_val,
prepare_response,
str_type,
str_val,
)
from google.api_core import grpc_helpers
from google.cloud.bigtable.data._sync_autogen._swappable_channel import SwappableChannel
CrossSync._Sync_Impl.add_mapping("grpc_helpers", grpc_helpers)
CrossSync._Sync_Impl.add_mapping("SwappableChannel", SwappableChannel)
@CrossSync._Sync_Impl.add_mapping_decorator("TestBigtableDataClient")
class TestBigtableDataClient:
@staticmethod
def _get_target_class():
return CrossSync._Sync_Impl.DataClient
@classmethod
def _make_client(cls, *args, use_emulator=True, **kwargs):
import os
env_mask = {}
if use_emulator:
env_mask["BIGTABLE_EMULATOR_HOST"] = "localhost"
import warnings
warnings.filterwarnings("ignore", category=RuntimeWarning)
else:
kwargs["credentials"] = kwargs.get("credentials", AnonymousCredentials())
kwargs["project"] = kwargs.get("project", "project-id")
with mock.patch.dict(os.environ, env_mask):
return cls._get_target_class()(*args, **kwargs)
def test_ctor(self):
expected_project = "project-id"
expected_credentials = AnonymousCredentials()
client = self._make_client(
project="project-id", credentials=expected_credentials, use_emulator=False
)
CrossSync._Sync_Impl.yield_to_event_loop()
assert client.project == expected_project
assert not client._active_instances
assert client._channel_refresh_task is not None
assert client.transport._credentials == expected_credentials
client.close()
def test_ctor_super_inits(self):
from google.cloud.client import ClientWithProject
from google.cloud.bigtable import __version__ as bigtable_version
from google.api_core import client_options as client_options_lib
from google.cloud.bigtable_v2.services.bigtable.transports.base import (
DEFAULT_CLIENT_INFO,
)
import copy
project = "project-id"
credentials = AnonymousCredentials()
client_options = {"api_endpoint": "foo.bar:1234"}
options_parsed = client_options_lib.from_dict(client_options)
with mock.patch.object(
CrossSync._Sync_Impl.GapicClient, "__init__"
) as bigtable_client_init:
bigtable_client_init.return_value = None
with mock.patch.object(
ClientWithProject, "__init__"
) as client_project_init:
client_project_init.return_value = None
try:
self._make_client(
project=project,
credentials=credentials,
client_options=options_parsed,
use_emulator=False,
)
except AttributeError:
pass
assert bigtable_client_init.call_count == 1
kwargs = bigtable_client_init.call_args[1]
assert kwargs["credentials"] == credentials
assert kwargs["client_options"] == options_parsed
expected_client_info = copy.copy(DEFAULT_CLIENT_INFO)
expected_client_info.client_library_version = (
f"{bigtable_version}-data"
if not CrossSync._Sync_Impl.is_async
else f"{bigtable_version}-data-async"
)
assert (
kwargs["client_info"].to_user_agent()
== expected_client_info.to_user_agent()
)
assert (
kwargs["client_info"].to_grpc_metadata()
== expected_client_info.to_grpc_metadata()
)
assert client_project_init.call_count == 1
kwargs = client_project_init.call_args[1]
assert kwargs["project"] == project
assert kwargs["credentials"] == credentials
assert kwargs["client_options"] == options_parsed
def test_ctor_legacy_client(self):
from google.api_core import client_options as client_options_lib
from google.api_core.gapic_v1.client_info import ClientInfo
project = "project-id"
credentials = AnonymousCredentials()
client_info = ClientInfo(gapic_version="1.2.3", user_agent="test-client-")
client_options = {"api_endpoint": "foo.bar:1234"}
options_parsed = client_options_lib.from_dict(client_options)
with mock.patch.object(
CrossSync._Sync_Impl.GapicClient, "__init__"
) as bigtable_client_init:
try:
client = self._make_client(
project=project,
credentials=credentials,
client_options=options_parsed,
use_emulator=False,
_client_info=client_info,
_disable_background_refresh=True,
)
assert client._disable_background_refresh
assert client.client_info is client_info
except TypeError:
pass
assert bigtable_client_init.call_count == 1
kwargs = bigtable_client_init.call_args[1]
assert kwargs["credentials"] == credentials
assert kwargs["client_options"] == options_parsed
def test_ctor_dict_options(self):
from google.api_core.client_options import ClientOptions
client_options = {"api_endpoint": "foo.bar:1234"}
with mock.patch.object(
CrossSync._Sync_Impl.GapicClient, "__init__"
) as bigtable_client_init:
try:
self._make_client(client_options=client_options)
except TypeError:
pass
bigtable_client_init.assert_called_once()
kwargs = bigtable_client_init.call_args[1]
called_options = kwargs["client_options"]
assert called_options.api_endpoint == "foo.bar:1234"
assert isinstance(called_options, ClientOptions)
with mock.patch.object(
self._get_target_class(), "_start_background_channel_refresh"
) as start_background_refresh:
client = self._make_client(
client_options=client_options, use_emulator=False
)
start_background_refresh.assert_called_once()
client.close()
def test_veneer_grpc_headers(self):
client_component = "data-async" if CrossSync._Sync_Impl.is_async else "data"
VENEER_HEADER_REGEX = re.compile(
"gapic\\/[0-9]+\\.[\\w.-]+ gax\\/[0-9]+\\.[\\w.-]+ gccl\\/[0-9]+\\.[\\w.-]+-"
+ client_component
+ " gl-python\\/[0-9]+\\.[\\w.-]+ grpc\\/[0-9]+\\.[\\w.-]+"
)
patch = mock.patch("google.api_core.gapic_v1.method.wrap_method")
with patch as gapic_mock:
client = self._make_client(project="project-id")
wrapped_call_list = gapic_mock.call_args_list
assert len(wrapped_call_list) > 0
for call in wrapped_call_list:
client_info = call.kwargs["client_info"]
assert client_info is not None, f"{call} has no client_info"
wrapped_user_agent_sorted = " ".join(
sorted(client_info.to_user_agent().split(" "))
)
assert VENEER_HEADER_REGEX.match(
wrapped_user_agent_sorted
), f"'{wrapped_user_agent_sorted}' does not match {VENEER_HEADER_REGEX}"
client.close()
def test__start_background_channel_refresh_task_exists(self):
client = self._make_client(project="project-id", use_emulator=False)
assert client._channel_refresh_task is not None
with mock.patch.object(asyncio, "create_task") as create_task:
client._start_background_channel_refresh()
create_task.assert_not_called()
client.close()
def test__start_background_channel_refresh(self):
client = self._make_client(project="project-id")
with mock.patch.object(
client, "_ping_and_warm_instances", CrossSync._Sync_Impl.Mock()
) as ping_and_warm:
client._emulator_host = None
client.transport._grpc_channel = CrossSync._Sync_Impl.SwappableChannel(
mock.Mock
)
client._start_background_channel_refresh()
assert client._channel_refresh_task is not None
assert isinstance(client._channel_refresh_task, CrossSync._Sync_Impl.Task)
CrossSync._Sync_Impl.sleep(0.1)
assert ping_and_warm.call_count == 1
client.close()
def test__start_background_channel_refresh_disable_background_refresh(self):
client = self._make_client(
project="project-id", _disable_background_refresh=True
)
with mock.patch.object(
client, "_ping_and_warm_instances", CrossSync._Sync_Impl.Mock()
) as ping_and_warm:
client._emulator_host = None
client.transport._grpc_channel = CrossSync._Sync_Impl.SwappableChannel(
mock.Mock
)
client._start_background_channel_refresh()
assert client._channel_refresh_task is None
ping_and_warm.assert_not_called()
def test__ping_and_warm_instances(self):
"""test ping and warm with mocked asyncio.gather"""
client_mock = mock.Mock()
client_mock._execute_ping_and_warms = (
lambda *args: self._get_target_class()._execute_ping_and_warms(
client_mock, *args
)
)
with mock.patch.object(
CrossSync._Sync_Impl, "gather_partials", CrossSync._Sync_Impl.Mock()
) as gather:
gather.side_effect = lambda partials, **kwargs: [None for _ in partials]
channel = mock.Mock()
client_mock._active_instances = []
result = self._get_target_class()._ping_and_warm_instances(
client_mock, channel=channel
)
assert len(result) == 0
assert gather.call_args[1]["return_exceptions"] is True
assert gather.call_args[1]["sync_executor"] == client_mock._executor
client_mock._active_instances = [(mock.Mock(), mock.Mock())] * 4
gather.reset_mock()
channel.reset_mock()
result = self._get_target_class()._ping_and_warm_instances(
client_mock, channel=channel
)
assert len(result) == 4
gather.assert_called_once()
partial_list = gather.call_args.args[0]
assert len(partial_list) == 4
grpc_call_args = channel.unary_unary().call_args_list
for idx, (_, kwargs) in enumerate(grpc_call_args):
(
expected_instance,
expected_app_profile,
) = client_mock._active_instances[idx]
request = kwargs["request"]
assert request["name"] == expected_instance
assert request["app_profile_id"] == expected_app_profile
metadata = kwargs["metadata"]
assert len(metadata) == 1
assert metadata[0][0] == "x-goog-request-params"
assert (
metadata[0][1]
== f"name={expected_instance}&app_profile_id={expected_app_profile}"
)
def test__ping_and_warm_single_instance(self):
"""should be able to call ping and warm with single instance"""
client_mock = mock.Mock()
client_mock._execute_ping_and_warms = (
lambda *args: self._get_target_class()._execute_ping_and_warms(
client_mock, *args
)
)
with mock.patch.object(
CrossSync._Sync_Impl, "gather_partials", CrossSync._Sync_Impl.Mock()
) as gather:
gather.side_effect = lambda *args, **kwargs: [fn() for fn in args[0]]
client_mock._active_instances = [mock.Mock()] * 100
test_key = ("test-instance", "test-app-profile")
result = self._get_target_class()._ping_and_warm_instances(
client_mock, test_key
)
assert len(result) == 1
grpc_call_args = (
client_mock.transport.grpc_channel.unary_unary().call_args_list
)
assert len(grpc_call_args) == 1
kwargs = grpc_call_args[0][1]
request = kwargs["request"]
assert request["name"] == "test-instance"
assert request["app_profile_id"] == "test-app-profile"
metadata = kwargs["metadata"]
assert len(metadata) == 1
assert metadata[0][0] == "x-goog-request-params"
assert (
metadata[0][1] == "name=test-instance&app_profile_id=test-app-profile"
)
@pytest.mark.parametrize(
"refresh_interval, wait_time, expected_sleep",
[(0, 0, 0), (0, 1, 0), (10, 0, 10), (10, 5, 5), (10, 10, 0), (10, 15, 0)],
)
def test__manage_channel_first_sleep(
self, refresh_interval, wait_time, expected_sleep
):
import time
with mock.patch.object(time, "monotonic") as monotonic:
monotonic.return_value = 0
with mock.patch.object(CrossSync._Sync_Impl, "event_wait") as sleep:
sleep.side_effect = asyncio.CancelledError
try:
client = self._make_client(project="project-id")
client._channel_init_time = -wait_time
client._manage_channel(refresh_interval, refresh_interval)
except asyncio.CancelledError:
pass
sleep.assert_called_once()
call_time = sleep.call_args[0][1]
assert (
abs(call_time - expected_sleep) < 0.1
), f"refresh_interval: {refresh_interval}, wait_time: {wait_time}, expected_sleep: {expected_sleep}"
client.close()
def test__manage_channel_ping_and_warm(self):
"""_manage channel should call ping and warm internally"""
import threading
client = self._make_client(project="project-id", use_emulator=True)
orig_channel = client.transport.grpc_channel
sleep_tuple = (
(asyncio, "sleep")
if CrossSync._Sync_Impl.is_async
else (threading.Event, "wait")
)
with mock.patch.object(*sleep_tuple) as sleep_mock:
sleep_mock.side_effect = [None, asyncio.CancelledError]
ping_and_warm = (
client._ping_and_warm_instances
) = CrossSync._Sync_Impl.Mock()
try:
client._manage_channel(10)
except asyncio.CancelledError:
pass
assert ping_and_warm.call_count == 2
assert client.transport.grpc_channel._channel != orig_channel
called_with = [call[1]["channel"] for call in ping_and_warm.call_args_list]
assert orig_channel in called_with
assert client.transport.grpc_channel._channel in called_with
@pytest.mark.parametrize(
"refresh_interval, num_cycles, expected_sleep",
[(None, 1, 60 * 35), (10, 10, 100), (10, 1, 10)],
)
def test__manage_channel_sleeps(self, refresh_interval, num_cycles, expected_sleep):
import time
import random
with mock.patch.object(random, "uniform") as uniform:
uniform.side_effect = lambda min_, max_: min_
with mock.patch.object(time, "time") as time_mock:
time_mock.return_value = 0
with mock.patch.object(CrossSync._Sync_Impl, "event_wait") as sleep:
sleep.side_effect = [None for i in range(num_cycles - 1)] + [
asyncio.CancelledError
]
client = self._make_client(project="project-id", use_emulator=True)
with mock.patch.object(
client.transport, "create_channel", CrossSync._Sync_Impl.Mock
):
try:
if refresh_interval is not None:
client._manage_channel(
refresh_interval, refresh_interval, grace_period=0
)
else:
client._manage_channel(grace_period=0)
except asyncio.CancelledError:
pass
assert sleep.call_count == num_cycles
total_sleep = sum([call[0][1] for call in sleep.call_args_list])
assert (
abs(total_sleep - expected_sleep) < 0.5
), f"refresh_interval={refresh_interval}, num_cycles={num_cycles}, expected_sleep={expected_sleep}"
client.close()
def test__manage_channel_random(self):
import random
with mock.patch.object(CrossSync._Sync_Impl, "event_wait") as sleep:
with mock.patch.object(random, "uniform") as uniform:
uniform.return_value = 0
try:
uniform.side_effect = asyncio.CancelledError
client = self._make_client(project="project-id")
except asyncio.CancelledError:
uniform.side_effect = None
uniform.reset_mock()
sleep.reset_mock()
with mock.patch.object(client.transport, "create_channel"):
min_val = 200
max_val = 205
uniform.side_effect = lambda min_, max_: min_
sleep.side_effect = [None, asyncio.CancelledError]
try:
client._manage_channel(min_val, max_val, grace_period=0)
except asyncio.CancelledError:
pass
assert uniform.call_count == 2
uniform_args = [call[0] for call in uniform.call_args_list]
for found_min, found_max in uniform_args:
assert found_min == min_val
assert found_max == max_val
@pytest.mark.parametrize("num_cycles", [0, 1, 10, 100])
def test__manage_channel_refresh(self, num_cycles):
expected_refresh = 0.5
grpc_lib = grpc.aio if CrossSync._Sync_Impl.is_async else grpc
new_channel = grpc_lib.insecure_channel("localhost:8080")
create_channel_mock = mock.Mock()
create_channel_mock.return_value = new_channel
refreshable_channel = CrossSync._Sync_Impl.SwappableChannel(create_channel_mock)
with mock.patch.object(CrossSync._Sync_Impl, "event_wait") as sleep:
sleep.side_effect = [None for i in range(num_cycles)] + [RuntimeError]
client = self._make_client(project="project-id")
client.transport._grpc_channel = refreshable_channel
create_channel_mock.reset_mock()
sleep.reset_mock()
try:
client._manage_channel(
refresh_interval_min=expected_refresh,
refresh_interval_max=expected_refresh,
grace_period=0,
)
except RuntimeError:
pass
assert sleep.call_count == num_cycles + 1
assert create_channel_mock.call_count == num_cycles
client.close()
def test__register_instance(self):
"""test instance registration"""
client_mock = mock.Mock()
client_mock._gapic_client.instance_path.side_effect = lambda a, b: f"prefix/{b}"
active_instances = set()
instance_owners = {}
client_mock._active_instances = active_instances
client_mock._instance_owners = instance_owners
client_mock._channel_refresh_task = None
client_mock._ping_and_warm_instances = CrossSync._Sync_Impl.Mock()
table_mock = mock.Mock()
self._get_target_class()._register_instance(
client_mock, "instance-1", table_mock.app_profile_id, id(table_mock)
)
assert client_mock._start_background_channel_refresh.call_count == 1
expected_key = ("prefix/instance-1", table_mock.app_profile_id)
assert len(active_instances) == 1
assert expected_key == tuple(list(active_instances)[0])
assert len(instance_owners) == 1
assert expected_key == tuple(list(instance_owners)[0])
client_mock._channel_refresh_task = mock.Mock()
table_mock2 = mock.Mock()
self._get_target_class()._register_instance(
client_mock, "instance-2", table_mock2.app_profile_id, id(table_mock2)
)
assert client_mock._start_background_channel_refresh.call_count == 1
assert (
client_mock._ping_and_warm_instances.call_args[0][0][0]
== "prefix/instance-2"
)
assert client_mock._ping_and_warm_instances.call_count == 1
assert len(active_instances) == 2
assert len(instance_owners) == 2
expected_key2 = ("prefix/instance-2", table_mock2.app_profile_id)
assert any(
[
expected_key2 == tuple(list(active_instances)[i])
for i in range(len(active_instances))
]
)
assert any(
[
expected_key2 == tuple(list(instance_owners)[i])
for i in range(len(instance_owners))
]
)
def test__register_instance_duplicate(self):
"""test double instance registration. Should be no-op"""
client_mock = mock.Mock()
client_mock._gapic_client.instance_path.side_effect = lambda a, b: f"prefix/{b}"
active_instances = set()
instance_owners = {}
client_mock._active_instances = active_instances
client_mock._instance_owners = instance_owners
client_mock._channel_refresh_task = object()
mock_channels = [mock.Mock()]
client_mock.transport.channels = mock_channels
client_mock._ping_and_warm_instances = CrossSync._Sync_Impl.Mock()
table_mock = mock.Mock()
expected_key = ("prefix/instance-1", table_mock.app_profile_id)
self._get_target_class()._register_instance(
client_mock, "instance-1", table_mock.app_profile_id, id(table_mock)
)
assert len(active_instances) == 1
assert expected_key == tuple(list(active_instances)[0])
assert len(instance_owners) == 1
assert expected_key == tuple(list(instance_owners)[0])
assert client_mock._ping_and_warm_instances.call_count == 1
self._get_target_class()._register_instance(
client_mock, "instance-1", table_mock.app_profile_id, id(table_mock)
)
assert len(active_instances) == 1
assert expected_key == tuple(list(active_instances)[0])
assert len(instance_owners) == 1
assert expected_key == tuple(list(instance_owners)[0])
assert client_mock._ping_and_warm_instances.call_count == 1
@pytest.mark.parametrize(
"insert_instances,expected_active,expected_owner_keys",
[
([("i", None)], [("i", None)], [("i", None)]),
([("i", "p")], [("i", "p")], [("i", "p")]),
([("1", "p"), ("1", "p")], [("1", "p")], [("1", "p")]),
(
[("1", "p"), ("2", "p")],
[("1", "p"), ("2", "p")],
[("1", "p"), ("2", "p")],
),
],
)
def test__register_instance_state(
self, insert_instances, expected_active, expected_owner_keys
):
"""test that active_instances and instance_owners are updated as expected"""
client_mock = mock.Mock()
client_mock._gapic_client.instance_path.side_effect = lambda a, b: b
active_instances = set()
instance_owners = {}
client_mock._active_instances = active_instances
client_mock._instance_owners = instance_owners
client_mock._channel_refresh_task = None
client_mock._ping_and_warm_instances = CrossSync._Sync_Impl.Mock()
table_mock = mock.Mock()
for instance, profile in insert_instances:
table_mock.app_profile_id = profile
self._get_target_class()._register_instance(
client_mock, instance, profile, id(table_mock)
)
assert len(active_instances) == len(expected_active)
assert len(instance_owners) == len(expected_owner_keys)
for expected in expected_active:
assert any(
[
expected == tuple(list(active_instances)[i])
for i in range(len(active_instances))
]
)
for expected in expected_owner_keys:
assert any(
[
expected == tuple(list(instance_owners)[i])
for i in range(len(instance_owners))
]
)
def test__remove_instance_registration(self):
client = self._make_client(project="project-id")
table = mock.Mock()
client._register_instance("instance-1", table.app_profile_id, id(table))
client._register_instance("instance-2", table.app_profile_id, id(table))
assert len(client._active_instances) == 2
assert len(client._instance_owners.keys()) == 2
instance_1_path = client._gapic_client.instance_path(
client.project, "instance-1"
)
instance_1_key = (instance_1_path, table.app_profile_id)
instance_2_path = client._gapic_client.instance_path(
client.project, "instance-2"
)
instance_2_key = (instance_2_path, table.app_profile_id)
assert len(client._instance_owners[instance_1_key]) == 1
assert list(client._instance_owners[instance_1_key])[0] == id(table)
assert len(client._instance_owners[instance_2_key]) == 1
assert list(client._instance_owners[instance_2_key])[0] == id(table)
success = client._remove_instance_registration(
"instance-1", table.app_profile_id, id(table)
)
assert success
assert len(client._active_instances) == 1
assert len(client._instance_owners[instance_1_key]) == 0
assert len(client._instance_owners[instance_2_key]) == 1
assert client._active_instances == {instance_2_key}
success = client._remove_instance_registration("fake-key", "profile", id(table))
assert not success
assert len(client._active_instances) == 1
client.close()
def test__multiple_table_registration(self):
"""registering with multiple tables with the same key should
add multiple owners to instance_owners, but only keep one copy
of shared key in active_instances"""
from google.cloud.bigtable.data._helpers import _WarmedInstanceKey
with self._make_client(project="project-id") as client:
with client.get_table("instance_1", "table_1") as table_1:
instance_1_path = client._gapic_client.instance_path(
client.project, "instance_1"
)
instance_1_key = _WarmedInstanceKey(
instance_1_path, table_1.app_profile_id
)
assert len(client._instance_owners[instance_1_key]) == 1
assert len(client._active_instances) == 1
assert id(table_1) in client._instance_owners[instance_1_key]
with client.get_table("instance_1", "table_2") as table_2:
assert table_2._register_instance_future is not None
table_2._register_instance_future.result()
assert len(client._instance_owners[instance_1_key]) == 2
assert len(client._active_instances) == 1
assert id(table_1) in client._instance_owners[instance_1_key]
assert id(table_2) in client._instance_owners[instance_1_key]
with client.get_table(
"instance_1", "table_3", app_profile_id="diff"
) as table_3:
assert table_3._register_instance_future is not None
table_3._register_instance_future.result()
instance_3_path = client._gapic_client.instance_path(
client.project, "instance_1"
)
instance_3_key = _WarmedInstanceKey(
instance_3_path, table_3.app_profile_id
)
assert len(client._instance_owners[instance_1_key]) == 2
assert len(client._instance_owners[instance_3_key]) == 1
assert len(client._active_instances) == 2
assert id(table_1) in client._instance_owners[instance_1_key]
assert id(table_2) in client._instance_owners[instance_1_key]
assert id(table_3) in client._instance_owners[instance_3_key]
assert len(client._active_instances) == 1
assert instance_1_key in client._active_instances
assert id(table_2) not in client._instance_owners[instance_1_key]
assert len(client._active_instances) == 0
assert instance_1_key not in client._active_instances
assert len(client._instance_owners[instance_1_key]) == 0
def test__multiple_instance_registration(self):
"""registering with multiple instance keys should update the key
in instance_owners and active_instances"""
from google.cloud.bigtable.data._helpers import _WarmedInstanceKey
with self._make_client(project="project-id") as client:
with client.get_table("instance_1", "table_1") as table_1:
assert table_1._register_instance_future is not None
table_1._register_instance_future.result()
with client.get_table("instance_2", "table_2") as table_2:
assert table_2._register_instance_future is not None
table_2._register_instance_future.result()
instance_1_path = client._gapic_client.instance_path(
client.project, "instance_1"
)
instance_1_key = _WarmedInstanceKey(
instance_1_path, table_1.app_profile_id
)
instance_2_path = client._gapic_client.instance_path(
client.project, "instance_2"
)
instance_2_key = _WarmedInstanceKey(
instance_2_path, table_2.app_profile_id
)
assert len(client._instance_owners[instance_1_key]) == 1
assert len(client._instance_owners[instance_2_key]) == 1
assert len(client._active_instances) == 2
assert id(table_1) in client._instance_owners[instance_1_key]
assert id(table_2) in client._instance_owners[instance_2_key]
assert len(client._active_instances) == 1
assert instance_1_key in client._active_instances
assert len(client._instance_owners[instance_2_key]) == 0
assert len(client._instance_owners[instance_1_key]) == 1
assert id(table_1) in client._instance_owners[instance_1_key]
assert len(client._active_instances) == 0
assert len(client._instance_owners[instance_1_key]) == 0
assert len(client._instance_owners[instance_2_key]) == 0
@pytest.mark.parametrize("method", ["get_table", "get_authorized_view"])
def test_get_api_surface(self, method):
"""test client.get_table and client.get_authorized_view"""
from google.cloud.bigtable.data._helpers import _WarmedInstanceKey
client = self._make_client(project="project-id")
assert not client._active_instances
expected_table_id = "table-id"
expected_instance_id = "instance-id"
expected_app_profile_id = "app-profile-id"
if method == "get_table":
surface = client.get_table(
expected_instance_id, expected_table_id, expected_app_profile_id
)
assert isinstance(
surface, CrossSync._Sync_Impl.TestTable._get_target_class()
)
elif method == "get_authorized_view":
surface = client.get_authorized_view(
expected_instance_id,
expected_table_id,
"view_id",
expected_app_profile_id,
)
assert isinstance(
surface, CrossSync._Sync_Impl.TestAuthorizedView._get_target_class()
)
assert (
surface.authorized_view_name
== f"projects/{client.project}/instances/{expected_instance_id}/tables/{expected_table_id}/authorizedViews/view_id"
)
else:
raise TypeError(f"unexpected method: {method}")
CrossSync._Sync_Impl.yield_to_event_loop()
assert surface.table_id == expected_table_id
assert (
surface.table_name
== f"projects/{client.project}/instances/{expected_instance_id}/tables/{expected_table_id}"
)
assert surface.instance_id == expected_instance_id
assert (
surface.instance_name
== f"projects/{client.project}/instances/{expected_instance_id}"
)
assert surface.app_profile_id == expected_app_profile_id
assert surface.client is client
instance_key = _WarmedInstanceKey(surface.instance_name, surface.app_profile_id)
assert instance_key in client._active_instances
assert client._instance_owners[instance_key] == {id(surface)}
client.close()
@pytest.mark.parametrize("method", ["get_table", "get_authorized_view"])
def test_api_surface_arg_passthrough(self, method):
"""All arguments passed in get_table and get_authorized_view should be sent to constructor"""
if method == "get_table":
surface_type = CrossSync._Sync_Impl.TestTable._get_target_class()
elif method == "get_authorized_view":
surface_type = CrossSync._Sync_Impl.TestAuthorizedView._get_target_class()
else:
raise TypeError(f"unexpected method: {method}")
with self._make_client(project="project-id") as client:
with mock.patch.object(surface_type, "__init__") as mock_constructor:
mock_constructor.return_value = None
assert not client._active_instances
expected_args = (
"table",
"instance",
"view",
"app_profile",
1,
"test",
{"test": 2},
)
expected_kwargs = {"hello": "world", "test": 2}
getattr(client, method)(*expected_args, **expected_kwargs)
mock_constructor.assert_called_once_with(
client, *expected_args, **expected_kwargs
)
@pytest.mark.parametrize("method", ["get_table", "get_authorized_view"])
def test_api_surface_context_manager(self, method):
"""get_table and get_authorized_view should work as context managers"""
from functools import partial
from google.cloud.bigtable.data._helpers import _WarmedInstanceKey
expected_table_id = "table-id"
expected_instance_id = "instance-id"
expected_app_profile_id = "app-profile-id"
expected_project_id = "project-id"
if method == "get_table":
surface_type = CrossSync._Sync_Impl.TestTable._get_target_class()
elif method == "get_authorized_view":
surface_type = CrossSync._Sync_Impl.TestAuthorizedView._get_target_class()
else:
raise TypeError(f"unexpected method: {method}")
with mock.patch.object(surface_type, "close") as close_mock:
with self._make_client(project=expected_project_id) as client:
if method == "get_table":
fn = partial(
client.get_table,
expected_instance_id,
expected_table_id,
expected_app_profile_id,
)
elif method == "get_authorized_view":
fn = partial(
client.get_authorized_view,
expected_instance_id,
expected_table_id,
"view_id",
expected_app_profile_id,
)
else:
raise TypeError(f"unexpected method: {method}")
with fn() as table:
CrossSync._Sync_Impl.yield_to_event_loop()
assert isinstance(table, surface_type)
assert table.table_id == expected_table_id
assert (
table.table_name
== f"projects/{expected_project_id}/instances/{expected_instance_id}/tables/{expected_table_id}"
)
assert table.instance_id == expected_instance_id
assert (
table.instance_name
== f"projects/{expected_project_id}/instances/{expected_instance_id}"
)
assert table.app_profile_id == expected_app_profile_id
assert table.client is client
instance_key = _WarmedInstanceKey(
table.instance_name, table.app_profile_id
)
assert instance_key in client._active_instances
assert client._instance_owners[instance_key] == {id(table)}
assert close_mock.call_count == 1
def test_close(self):
client = self._make_client(project="project-id", use_emulator=False)
task = client._channel_refresh_task
assert task is not None
assert not task.done()
with mock.patch.object(
client.transport, "close", CrossSync._Sync_Impl.Mock()
) as close_mock:
client.close()
close_mock.assert_called_once()
assert task.done()
assert client._channel_refresh_task is None
def test_close_with_timeout(self):
expected_timeout = 19
client = self._make_client(project="project-id", use_emulator=False)
with mock.patch.object(
CrossSync._Sync_Impl, "wait", CrossSync._Sync_Impl.Mock()
) as wait_for_mock:
client.close(timeout=expected_timeout)
wait_for_mock.assert_called_once()
assert wait_for_mock.call_args[1]["timeout"] == expected_timeout
client.close()
def test_context_manager(self):
from functools import partial
close_mock = CrossSync._Sync_Impl.Mock()
true_close = None
with self._make_client(project="project-id", use_emulator=False) as client:
true_close = partial(client.close)
client.close = close_mock
assert not client._channel_refresh_task.done()
assert client.project == "project-id"
assert client._active_instances == set()
close_mock.assert_not_called()
close_mock.assert_called_once()
true_close()
def test_default_universe_domain(self):
"""When not passed, universe_domain should default to googleapis.com"""
with self._make_client(project="project-id", credentials=None) as client:
assert client.universe_domain == "googleapis.com"
assert client.api_endpoint == "bigtable.googleapis.com"
def test_custom_universe_domain(self):
"""test with a customized universe domain value and emulator enabled"""
universe_domain = "test-universe.test"
options = client_options.ClientOptions(universe_domain=universe_domain)
with self._make_client(
project="project_id",
client_options=options,
use_emulator=True,
credentials=None,
) as client:
assert client.universe_domain == universe_domain
assert client.api_endpoint == f"bigtable.{universe_domain}"
def test_configured_universe_domain_matches_GDU(self):
"""that configured universe domain succeeds with matched GDU credentials."""
universe_domain = "googleapis.com"
options = client_options.ClientOptions(universe_domain=universe_domain)
with self._make_client(
project="project_id", client_options=options, credentials=None
) as client:
assert client.universe_domain == "googleapis.com"
assert client.api_endpoint == "bigtable.googleapis.com"
def test_credential_universe_domain_matches_GDU(self):
"""Test with credentials"""
creds = AnonymousCredentials()
creds._universe_domain = "googleapis.com"
with self._make_client(project="project_id", credentials=creds) as client:
assert client.universe_domain == "googleapis.com"
assert client.api_endpoint == "bigtable.googleapis.com"
def test_anomynous_credential_universe_domain(self):
"""Anomynopus credentials should use default universe domain"""
creds = AnonymousCredentials()
with self._make_client(project="project_id", credentials=creds) as client:
assert client.universe_domain == "googleapis.com"
assert client.api_endpoint == "bigtable.googleapis.com"
def test_configured_universe_domain_mismatched_credentials(self):
"""Test that configured universe domain errors with mismatched universe
domain credentials."""
universe_domain = "test-universe.test"
options = client_options.ClientOptions(universe_domain=universe_domain)
creds = AnonymousCredentials()
creds._universe_domain = "different-universe"
with pytest.raises(ValueError) as exc:
self._make_client(
project="project_id",
client_options=options,
use_emulator=False,
credentials=creds,
)
err_msg = f"The configured universe domain ({universe_domain}) does not match the universe domain found in the credentials ({creds.universe_domain}). If you haven't configured the universe domain explicitly, `googleapis.com` is the default."
assert exc.value.args[0] == err_msg
def test_configured_universe_domain_matches_credentials(self):
"""Test that configured universe domain succeeds with matching universe
domain credentials."""
universe_domain = "test-universe.test"
options = client_options.ClientOptions(universe_domain=universe_domain)
creds = AnonymousCredentials()
creds._universe_domain = universe_domain
with self._make_client(
project="project_id", credentials=creds, client_options=options
) as client:
assert client.universe_domain == universe_domain
assert client.api_endpoint == f"bigtable.{universe_domain}"
@CrossSync._Sync_Impl.add_mapping_decorator("TestTable")
class TestTable:
def _make_client(self, *args, **kwargs):
return CrossSync._Sync_Impl.TestBigtableDataClient._make_client(*args, **kwargs)
@staticmethod
def _get_target_class():
return CrossSync._Sync_Impl.Table
def _make_one(
self,
client,
instance_id="instance",
table_id="table",
app_profile_id=None,
**kwargs,
):
return self._get_target_class()(
client, instance_id, table_id, app_profile_id, **kwargs
)
def test_ctor(self):
from google.cloud.bigtable.data._helpers import _WarmedInstanceKey