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
|
# 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(location1_endpoint, "Read") is False
assert lc.is_endpoint_unavailable(location1_endpoint, "None") is False
assert lc.is_endpoint_unavailable(location1_endpoint, "Write") is False
lc.mark_endpoint_unavailable_for_read(location1_endpoint, False)
assert lc.is_endpoint_unavailable(location1_endpoint, "Read")
assert lc.is_endpoint_unavailable(location1_endpoint, "None") is False
assert lc.is_endpoint_unavailable(location1_endpoint, "Write") is False
lc.mark_endpoint_unavailable_for_write(location2_endpoint, False)
assert lc.is_endpoint_unavailable(location2_endpoint, "Read") is False
assert lc.is_endpoint_unavailable(location2_endpoint, "None") is False
assert lc.is_endpoint_unavailable(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) == 2
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) == 3
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 == region.get_primary():
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 == region.get_primary():
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)
# With the updated logic, we should retry the unavailable preferred locations
# instead of falling back to the default endpoint.
assert read_resolved == location1_endpoint
assert write_resolved == location1_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
def test_resolve_endpoint_unavailable_and_excluded_preferred_regions(self):
# Scenario: All preferred read locations are unavailable AND in the excluded list.
# Expected: Fallback to the primary write region.
connection_policy = documents.ConnectionPolicy()
connection_policy.ExcludedLocations = [location1_name, location4_name]
lc = refresh_location_cache([location1_name, location4_name], True, connection_policy)
db_acc = create_database_account(True)
lc.perform_on_database_account_read(db_acc)
# Mark all preferred read locations as unavailable
lc.mark_endpoint_unavailable_for_read(location1_endpoint, True)
lc.mark_endpoint_unavailable_for_read(location4_endpoint, True)
# Create a read request
read_doc_request = RequestObject(ResourceType.Document, _OperationType.Read, None)
# Resolve the endpoint for the read request
read_doc_resolved = lc.resolve_service_endpoint(read_doc_request)
# All preferred read locations ([loc1, loc4]) are excluded.
# The fallback for read is the primary write region, which is loc1.
assert read_doc_resolved == location1_endpoint
# Scenario: All preferred write locations are unavailable AND in the excluded list.
# Expected: Fallback to the default endpoint.
connection_policy.ExcludedLocations = [location1_name, location2_name]
lc = refresh_location_cache([location1_name, location2_name], True, connection_policy)
db_acc = create_database_account(True)
lc.perform_on_database_account_read(db_acc)
# Mark preferred write locations as unavailable
lc.mark_endpoint_unavailable_for_write(location1_endpoint, True)
lc.mark_endpoint_unavailable_for_write(location2_endpoint, True)
# Create a write request
write_doc_request = RequestObject(ResourceType.Document, _OperationType.Create, None)
# Resolve the endpoint for the write request
write_doc_resolved = lc.resolve_service_endpoint(write_doc_request)
# All preferred write locations ([loc1, loc2]) are excluded.
# The fallback for write is the default_endpoint.
assert write_doc_resolved == default_endpoint
def test_resolve_endpoint_unavailable_and_excluded_on_request(self):
# Scenario: All preferred read locations are unavailable AND in the excluded list on the request.
# Expected: Fallback to the primary write region.
lc = refresh_location_cache([location1_name, location4_name], True)
db_acc = create_database_account(True)
lc.perform_on_database_account_read(db_acc)
# Mark all preferred read locations as unavailable
lc.mark_endpoint_unavailable_for_read(location1_endpoint, True)
lc.mark_endpoint_unavailable_for_read(location4_endpoint, True)
# Create a read request and set excluded locations
read_doc_request = RequestObject(ResourceType.Document, _OperationType.Read, None)
read_doc_request.excluded_locations = [location1_name, location4_name]
# Resolve the endpoint for the read request
read_doc_resolved = lc.resolve_service_endpoint(read_doc_request)
# All preferred read locations ([loc1, loc4]) are excluded.
# The fallback for read is the primary write region, which is loc1.
assert read_doc_resolved == location1_endpoint
# Scenario: All preferred write locations are unavailable AND in the excluded list on the request.
# Expected: Fallback to the default endpoint.
lc = refresh_location_cache([location1_name, location2_name], True)
db_acc = create_database_account(True)
lc.perform_on_database_account_read(db_acc)
# Mark preferred write locations as unavailable
lc.mark_endpoint_unavailable_for_write(location1_endpoint, True)
lc.mark_endpoint_unavailable_for_write(location2_endpoint, True)
# Create a write request and set excluded locations
write_doc_request = RequestObject(ResourceType.Document, _OperationType.Create, None)
write_doc_request.excluded_locations = [location1_name, location2_name]
# Resolve the endpoint for the write request
write_doc_resolved = lc.resolve_service_endpoint(write_doc_request)
# All preferred write locations ([loc1, loc2]) are excluded.
# The fallback for write is the default_endpoint.
assert write_doc_resolved == default_endpoint
if __name__ == "__main__":
unittest.main()
|
