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
|
# The MIT License (MIT)
# Copyright (c) Microsoft Corporation. All rights reserved.
import time
import unittest
from typing import Mapping, Any
import pytest
from azure.cosmos import documents
from azure.cosmos.documents import DatabaseAccount, _OperationType
from azure.cosmos.http_constants import ResourceType
from azure.cosmos._location_cache import LocationCache
from azure.cosmos._request_object import RequestObject
default_endpoint = "https://default.documents.azure.com"
location1_name = "location1"
location2_name = "location2"
location3_name = "location3"
location4_name = "location4"
location1_endpoint = "https://location1.documents.azure.com"
location2_endpoint = "https://location2.documents.azure.com"
location3_endpoint = "https://location3.documents.azure.com"
location4_endpoint = "https://location4.documents.azure.com"
refresh_time_interval_in_ms = 1000
def create_database_account(enable_multiple_writable_locations):
db_acc = DatabaseAccount()
db_acc._WritableLocations = [{"name": location1_name, "databaseAccountEndpoint": location1_endpoint},
{"name": location2_name, "databaseAccountEndpoint": location2_endpoint},
{"name": location3_name, "databaseAccountEndpoint": location3_endpoint}]
db_acc._ReadableLocations = [{"name": location1_name, "databaseAccountEndpoint": location1_endpoint},
{"name": location2_name, "databaseAccountEndpoint": location2_endpoint},
{"name": location4_name, "databaseAccountEndpoint": location4_endpoint}]
db_acc._EnableMultipleWritableLocations = enable_multiple_writable_locations
return db_acc
def refresh_location_cache(preferred_locations, use_multiple_write_locations, connection_policy=documents.ConnectionPolicy()):
connection_policy.PreferredLocations = preferred_locations
connection_policy.UseMultipleWriteLocations = use_multiple_write_locations
lc = LocationCache(default_endpoint=default_endpoint,
connection_policy=connection_policy)
return lc
@pytest.mark.cosmosEmulator
class TestLocationCache:
def test_mark_endpoint_unavailable(self):
lc = refresh_location_cache([], False)
# mark unavailable for read
lc.mark_endpoint_unavailable_for_read(location1_endpoint, True)
location1_info = lc.location_unavailability_info_by_endpoint[location1_endpoint]
assert location1_info['operationType'] == {'Read'}
# mark unavailable for write
time.sleep(1)
lc.mark_endpoint_unavailable_for_write(location1_endpoint, False)
location1_info = lc.location_unavailability_info_by_endpoint[location1_endpoint]
assert location1_info['operationType'] == {'Read', 'Write'}
def test_is_endpoint_unavailable(self):
lc = refresh_location_cache([], False)
assert lc.is_endpoint_unavailable_internal(location1_endpoint, "Read") is False
assert lc.is_endpoint_unavailable_internal(location1_endpoint, "None") is False
assert lc.is_endpoint_unavailable_internal(location1_endpoint, "Write") is False
lc.mark_endpoint_unavailable_for_read(location1_endpoint, False)
assert lc.is_endpoint_unavailable_internal(location1_endpoint, "Read")
assert lc.is_endpoint_unavailable_internal(location1_endpoint, "None") is False
assert lc.is_endpoint_unavailable_internal(location1_endpoint, "Write") is False
lc.mark_endpoint_unavailable_for_write(location2_endpoint, False)
assert lc.is_endpoint_unavailable_internal(location2_endpoint, "Read") is False
assert lc.is_endpoint_unavailable_internal(location2_endpoint, "None") is False
assert lc.is_endpoint_unavailable_internal(location2_endpoint, "Write")
location1_info = lc.location_unavailability_info_by_endpoint[location1_endpoint]
lc.location_unavailability_info_by_endpoint[location1_endpoint] = location1_info
def test_endpoints_to_health_check(self):
lc = refresh_location_cache([location4_name], False)
db_acc = create_database_account(False)
lc.perform_on_database_account_read(db_acc)
# check endpoints to health check
endpoints = lc.endpoints_to_health_check()
assert len(endpoints) == 3
assert default_endpoint in endpoints
assert location1_endpoint in endpoints
assert location4_endpoint in endpoints
def test_get_locations(self):
lc = refresh_location_cache([], False)
db_acc = create_database_account(False)
lc.perform_on_database_account_read(db_acc)
# check read endpoints without preferred locations
read_regions = lc.get_read_regional_routing_contexts()
assert len(read_regions) == 1
assert read_regions[0].get_primary() == location1_endpoint
# check read endpoints with preferred locations
lc = refresh_location_cache([location1_name, location2_name, location4_name], False)
lc.perform_on_database_account_read(db_acc)
read_regions = lc.get_read_regional_routing_contexts()
assert len(read_regions) == len(db_acc.ReadableLocations)
for read_region in db_acc.ReadableLocations:
found_endpoint = False
endpoint = read_region['databaseAccountEndpoint']
for region in read_regions:
if endpoint in (region.get_primary(), region.get_alternate()):
found_endpoint = True
assert found_endpoint
# check write endpoints
write_regions = lc.get_write_regional_routing_contexts()
assert len(write_regions) == len(db_acc.WritableLocations)
for write_region in db_acc.WritableLocations:
found_endpoint = False
endpoint = write_region['databaseAccountEndpoint']
for region in write_regions:
if endpoint in (region.get_primary(), region.get_alternate()):
found_endpoint = True
assert found_endpoint
def test_resolve_request_endpoint_preferred_regions(self):
lc = refresh_location_cache([location1_name, location3_name, location4_name], True)
db_acc = create_database_account(True)
lc.perform_on_database_account_read(db_acc)
write_doc_request = RequestObject(ResourceType.Document, _OperationType.Create, None)
read_doc_request = RequestObject(ResourceType.Document, _OperationType.Read, None)
# resolve both document requests with all regions available
write_doc_resolved = lc.resolve_service_endpoint(write_doc_request)
read_doc_resolved = lc.resolve_service_endpoint(read_doc_request)
assert write_doc_resolved == read_doc_resolved
# mark main region unavailable and try again
lc.mark_endpoint_unavailable_for_read(location1_endpoint, True)
lc.mark_endpoint_unavailable_for_write(location1_endpoint, True)
read_doc_resolved = lc.resolve_service_endpoint(read_doc_request)
write_doc_resolved = lc.resolve_service_endpoint(write_doc_request)
assert read_doc_resolved == location4_endpoint
assert write_doc_resolved == location3_endpoint
# mark next preferred region as unavailable - no preferred endpoints left
lc.mark_endpoint_unavailable_for_read(location4_endpoint, True)
lc.mark_endpoint_unavailable_for_write(location3_endpoint, True)
read_resolved = lc.resolve_service_endpoint(read_doc_request)
write_resolved = lc.resolve_service_endpoint(write_doc_request)
assert read_resolved == write_resolved
assert read_resolved == default_endpoint
@pytest.mark.parametrize("test_type",["OnClient", "OnRequest", "OnBoth"])
def test_get_applicable_regional_endpoints_excluded_regions(self, test_type):
# Init test data
if test_type == "OnClient":
excluded_locations_on_client_list = [
[location1_name],
[location1_name, location2_name],
[location1_name, location2_name, location3_name],
[location4_name],
[],
]
excluded_locations_on_requests_list = [None] * 5
elif test_type == "OnRequest":
excluded_locations_on_client_list = [[]] * 5
excluded_locations_on_requests_list = [
[location1_name],
[location1_name, location2_name],
[location1_name, location2_name, location3_name],
[location4_name],
[],
]
else:
excluded_locations_on_client_list = [
[location1_name],
[location1_name, location2_name, location3_name],
[location1_name, location2_name],
[location2_name],
[location1_name, location2_name, location3_name],
]
excluded_locations_on_requests_list = [
[location1_name],
[location1_name, location2_name],
[location1_name, location2_name, location3_name],
[location4_name],
[],
]
expected_read_endpoints_list = [
[location2_endpoint],
[location1_endpoint],
[location1_endpoint],
[location1_endpoint, location2_endpoint],
[location1_endpoint, location2_endpoint],
]
expected_write_endpoints_list = [
[location2_endpoint, location3_endpoint],
[location3_endpoint],
[default_endpoint],
[location1_endpoint, location2_endpoint, location3_endpoint],
[location1_endpoint, location2_endpoint, location3_endpoint],
]
# Loop over each test cases
for excluded_locations_on_client, excluded_locations_on_requests, expected_read_endpoints, expected_write_endpoints in zip(excluded_locations_on_client_list, excluded_locations_on_requests_list, expected_read_endpoints_list, expected_write_endpoints_list):
# Init excluded_locations in ConnectionPolicy
connection_policy = documents.ConnectionPolicy()
connection_policy.ExcludedLocations = excluded_locations_on_client
# Init location_cache
location_cache = refresh_location_cache([location1_name, location2_name, location3_name], True,
connection_policy)
database_account = create_database_account(True)
location_cache.perform_on_database_account_read(database_account)
# Init requests and set excluded regions on requests
write_doc_request = RequestObject(ResourceType.Document, _OperationType.Create, None)
write_doc_request.excluded_locations = excluded_locations_on_requests
read_doc_request = RequestObject(ResourceType.Document, _OperationType.Read, None)
read_doc_request.excluded_locations = excluded_locations_on_requests
# Test if read endpoints were correctly filtered on client level
read_regional_routing_contexts = location_cache._get_applicable_read_regional_routing_contexts(read_doc_request)
read_doc_endpoint = [regional_routing_contexts.get_primary() for regional_routing_contexts in read_regional_routing_contexts]
assert read_doc_endpoint == expected_read_endpoints
# Test if write endpoints were correctly filtered on client level
write_regional_routing_contexts = location_cache._get_applicable_write_regional_routing_contexts(write_doc_request)
write_doc_endpoint = [regional_routing_contexts.get_primary() for regional_routing_contexts in write_regional_routing_contexts]
assert write_doc_endpoint == expected_write_endpoints
def test_set_excluded_locations_for_requests(self):
# Init excluded_locations in ConnectionPolicy
excluded_locations_on_client = [location1_name, location2_name]
connection_policy = documents.ConnectionPolicy()
connection_policy.ExcludedLocations = excluded_locations_on_client
# Init location_cache
location_cache = refresh_location_cache([location1_name, location2_name, location3_name], True,
connection_policy)
database_account = create_database_account(True)
location_cache.perform_on_database_account_read(database_account)
# Test setting excluded locations
excluded_locations = [location1_name]
options: Mapping[str, Any] = {"excludedLocations": excluded_locations}
expected_excluded_locations = excluded_locations
read_doc_request = RequestObject(ResourceType.Document, _OperationType.Create, None)
read_doc_request.set_excluded_location_from_options(options)
actual_excluded_locations = read_doc_request.excluded_locations
assert actual_excluded_locations == expected_excluded_locations
expected_read_endpoints = [location2_endpoint]
read_regional_routing_contexts = location_cache._get_applicable_read_regional_routing_contexts(read_doc_request)
read_doc_endpoint = [regional_routing_contexts.get_primary() for regional_routing_contexts in read_regional_routing_contexts]
assert read_doc_endpoint == expected_read_endpoints
# Test setting excluded locations with None value
expected_error_message = ("Excluded locations cannot be None. "
"If you want to remove all excluded locations, try passing an empty list.")
with pytest.raises(ValueError) as e:
options: Mapping[str, Any] = {"excludedLocations": None}
doc_request = RequestObject(ResourceType.Document, _OperationType.Create, None)
doc_request.set_excluded_location_from_options(options)
assert str(
e.value) == expected_error_message
if __name__ == "__main__":
unittest.main()
|
