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
|
# The MIT License (MIT)
# Copyright (c) Microsoft Corporation. All rights reserved.
import asyncio
import time
import unittest
import uuid
from typing import List
import pytest
import pytest_asyncio
from azure.core.exceptions import ServiceRequestError
import test_config
from azure.cosmos import DatabaseAccount, _location_cache
from azure.cosmos._location_cache import RegionalRoutingContext
from azure.cosmos.aio import CosmosClient, _global_endpoint_manager_async, _cosmos_client_connection_async, \
_retry_utility_async
COLLECTION = "created_collection"
REGION_1 = "East US"
REGION_2 = "West US"
REGIONS = [REGION_1, REGION_2]
@pytest_asyncio.fixture()
async def setup():
if (TestHealthCheckAsync.masterKey == '[YOUR_KEY_HERE]' or
TestHealthCheckAsync.host == '[YOUR_ENDPOINT_HERE]'):
raise Exception(
"You must specify your Azure Cosmos account values for "
"'masterKey' and 'host' at the top of this class to run the "
"tests.")
client = CosmosClient(TestHealthCheckAsync.host, TestHealthCheckAsync.masterKey, consistency_level="Session")
created_database = client.get_database_client(TestHealthCheckAsync.TEST_DATABASE_ID)
created_collection = created_database.get_container_client(TestHealthCheckAsync.TEST_CONTAINER_SINGLE_PARTITION_ID)
yield {
COLLECTION: created_collection
}
await client.close()
@pytest.mark.cosmosEmulator
@pytest.mark.asyncio
@pytest.mark.usefixtures("setup")
class TestHealthCheckAsync:
host = test_config.TestConfig.host
masterKey = test_config.TestConfig.masterKey
connectionPolicy = test_config.TestConfig.connectionPolicy
TEST_DATABASE_ID = test_config.TestConfig.TEST_DATABASE_ID
TEST_CONTAINER_SINGLE_PARTITION_ID = test_config.TestConfig.TEST_SINGLE_PARTITION_CONTAINER_ID
async def test_health_check_success_startup_async(self, setup):
# checks at startup that we perform a health check on all the necessary endpoints
self.original_getDatabaseAccountStub = _global_endpoint_manager_async._GlobalEndpointManager._GetDatabaseAccountStub
self.original_getDatabaseAccountCheck = _cosmos_client_connection_async.CosmosClientConnection._GetDatabaseAccountCheck
mock_get_database_account_check = self.MockGetDatabaseAccountCheck()
_global_endpoint_manager_async._GlobalEndpointManager._GetDatabaseAccountStub = (
self.MockGetDatabaseAccount(REGIONS))
_cosmos_client_connection_async.CosmosClientConnection._GetDatabaseAccountCheck = mock_get_database_account_check
try:
client = CosmosClient(self.host, self.masterKey, preferred_locations=REGIONS)
# this will setup the location cache
client.client_connection._global_endpoint_manager.refresh_needed = True
await client.client_connection._global_endpoint_manager.refresh_endpoint_list(None)
finally:
_global_endpoint_manager_async._GlobalEndpointManager._GetDatabaseAccountStub = self.original_getDatabaseAccountStub
_cosmos_client_connection_async.CosmosClientConnection._GetDatabaseAccountCheck = self.original_getDatabaseAccountCheck
expected_regional_routing_context = []
locational_endpoint = _location_cache.LocationCache.GetLocationalEndpoint(self.host, REGION_1)
assert mock_get_database_account_check.counter == 2
endpoint = locational_endpoint
expected_regional_routing_context.append(RegionalRoutingContext(endpoint))
locational_endpoint = _location_cache.LocationCache.GetLocationalEndpoint(self.host, REGION_2)
expected_regional_routing_context.append(RegionalRoutingContext(locational_endpoint))
read_regional_routing_context = client.client_connection._global_endpoint_manager.location_cache.read_regional_routing_contexts
assert read_regional_routing_context == expected_regional_routing_context
await client.close()
async def test_health_check_failure_startup_async(self, setup):
# checks at startup that the health check will mark endpoints as unavailable if it gets an error
self.original_getDatabaseAccountStub = _global_endpoint_manager_async._GlobalEndpointManager._GetDatabaseAccountStub
_global_endpoint_manager_async._GlobalEndpointManager._GetDatabaseAccountStub = (
self.MockGetDatabaseAccount(REGIONS))
# don't mock database account check because we want it to fail and the emulator doesn't have extra regions
try:
client = CosmosClient(self.host, self.masterKey, preferred_locations=REGIONS)
# this will setup the location cache
client.client_connection._global_endpoint_manager.refresh_needed = True
await client.client_connection._global_endpoint_manager.refresh_endpoint_list(None)
finally:
_global_endpoint_manager_async._GlobalEndpointManager._GetDatabaseAccountStub = self.original_getDatabaseAccountStub
expected_endpoints = []
for region in REGIONS:
locational_endpoint = _location_cache.LocationCache.GetLocationalEndpoint(self.host, region)
expected_endpoints.append(locational_endpoint)
unavailable_endpoint_info = client.client_connection._global_endpoint_manager.location_cache.location_unavailability_info_by_endpoint
assert len(unavailable_endpoint_info) == len(expected_endpoints)
for expected_endpoint in expected_endpoints:
assert expected_endpoint in unavailable_endpoint_info.keys()
await client.close()
async def test_health_check_background(self, setup):
# makes sure the health check is in the background and doesn't block by mocking it with a large sleep value
self.original_health_check = _global_endpoint_manager_async._GlobalEndpointManager._endpoints_health_check
_global_endpoint_manager_async._GlobalEndpointManager._endpoints_health_check = self.mock_health_check
start_time = time.time()
try:
setup[COLLECTION].client_connection._global_endpoint_manager.startup = False
for i in range(5):
await setup[COLLECTION].create_item(body={'id': 'item' + str(uuid.uuid4()), 'pk': 'pk'})
finally:
_global_endpoint_manager_async._GlobalEndpointManager._endpoints_health_check = self.original_health_check
end_time = time.time()
duration = end_time - start_time
assert duration < 2, f"Test took too long: {duration} seconds"
async def test_health_check_background_fail(self, setup):
# makes sure exceptions in the health check aren't bubbled up but swallowed
# by mocking health check with an error
self.original_health_check = _global_endpoint_manager_async._GlobalEndpointManager._endpoints_health_check
_global_endpoint_manager_async._GlobalEndpointManager._endpoints_health_check = self.mock_health_check_failure
try:
setup[COLLECTION].client_connection._global_endpoint_manager.startup = False
for i in range(20):
await setup[COLLECTION].create_item(body={'id': 'item' + str(uuid.uuid4()), 'pk': 'pk'})
finally:
_global_endpoint_manager_async._GlobalEndpointManager._endpoints_health_check = self.original_health_check
async def test_health_check_success_async(self, setup):
# checks the background health check works as expected when all endpoints healthy
self.original_getDatabaseAccountStub = _global_endpoint_manager_async._GlobalEndpointManager._GetDatabaseAccountStub
self.original_getDatabaseAccountCheck = _cosmos_client_connection_async.CosmosClientConnection._GetDatabaseAccountCheck
self.original_preferred_locations = setup[COLLECTION].client_connection.connection_policy.PreferredLocations
setup[COLLECTION].client_connection.connection_policy.PreferredLocations = REGIONS
mock_get_database_account_check = self.MockGetDatabaseAccountCheck()
_global_endpoint_manager_async._GlobalEndpointManager._GetDatabaseAccountStub = (
self.MockGetDatabaseAccount(REGIONS))
self.OriginalExecuteFunction = _retry_utility_async.ExecuteFunctionAsync
_cosmos_client_connection_async.CosmosClientConnection._GetDatabaseAccountCheck = mock_get_database_account_check
async def mock_execute_function(function, *args, **kwargs):
if args:
args[4].url = args[4].url.replace('-eastus', '').replace('-westus', '')
args[1].location_endpoint_to_route = args[1].location_endpoint_to_route.replace('-eastus', '').replace('-westus', '')
if args[1].endpoint_override:
args[1].endpoint_override = args[1].endpoint_override.replace('-eastus', '').replace('-westus', '')
return await self.OriginalExecuteFunction(function, *args, **kwargs)
_retry_utility_async.ExecuteFunctionAsync = mock_execute_function
try:
setup[COLLECTION].client_connection._global_endpoint_manager.startup = False
setup[COLLECTION].client_connection._global_endpoint_manager.refresh_needed = True
for i in range(2):
await setup[COLLECTION].create_item(body={'id': 'item' + str(uuid.uuid4()), 'pk': 'pk'})
finally:
_global_endpoint_manager_async._GlobalEndpointManager._GetDatabaseAccountStub = self.original_getDatabaseAccountStub
_cosmos_client_connection_async.CosmosClientConnection._GetDatabaseAccountCheck = self.original_getDatabaseAccountCheck
setup[COLLECTION].client_connection.connection_policy.PreferredLocations = self.original_preferred_locations
_retry_utility_async.ExecuteFunctionAsync = self.OriginalExecuteFunction
expected_regional_routing_contexts = []
locational_endpoint = _location_cache.LocationCache.GetLocationalEndpoint(self.host, REGION_1)
assert mock_get_database_account_check.counter == 2
expected_regional_routing_contexts.append(RegionalRoutingContext(locational_endpoint))
locational_endpoint = _location_cache.LocationCache.GetLocationalEndpoint(self.host, REGION_2)
expected_regional_routing_contexts.append(RegionalRoutingContext(locational_endpoint))
read_regional_routing_context = setup[COLLECTION].client_connection._global_endpoint_manager.location_cache.read_regional_routing_contexts
assert read_regional_routing_context == expected_regional_routing_contexts
async def test_health_check_failure_async(self, setup):
# checks the background health check works as expected when all endpoints unhealthy - it should mark the endpoints unavailable
setup[COLLECTION].client_connection._global_endpoint_manager.location_cache.location_unavailability_info_by_endpoint.clear()
self.original_getDatabaseAccountStub = _global_endpoint_manager_async._GlobalEndpointManager._GetDatabaseAccountStub
_global_endpoint_manager_async._GlobalEndpointManager._GetDatabaseAccountStub = (
self.MockGetDatabaseAccount(REGIONS))
self.original_getDatabaseAccountCheck = _cosmos_client_connection_async.CosmosClientConnection._GetDatabaseAccountCheck
self.original_preferred_locations = setup[COLLECTION].client_connection.connection_policy.PreferredLocations
self.OriginalExecuteFunction = _retry_utility_async.ExecuteFunctionAsync
mock_get_database_account_check = self.MockGetDatabaseAccountCheckError()
_cosmos_client_connection_async.CosmosClientConnection._GetDatabaseAccountCheck = mock_get_database_account_check
setup[COLLECTION].client_connection.connection_policy.PreferredLocations = REGIONS
async def mock_execute_function(function, *args, **kwargs):
if args:
args[4].url = args[4].url.replace('-eastus', '').replace('-westus', '')
args[1].location_endpoint_to_route = args[1].location_endpoint_to_route.replace('-eastus', '').replace('-westus', '')
if args[1].endpoint_override:
args[1].endpoint_override = args[1].endpoint_override.replace('-eastus', '').replace('-westus', '')
return await self.OriginalExecuteFunction(function, *args, **kwargs)
_retry_utility_async.ExecuteFunctionAsync = mock_execute_function
try:
setup[COLLECTION].client_connection._global_endpoint_manager.startup = False
setup[COLLECTION].client_connection._global_endpoint_manager.refresh_needed = True
# Trigger the background health check
await setup[COLLECTION].create_item(body={'id': 'item' + str(uuid.uuid4()), 'pk': 'pk'})
# Poll until the background task marks the endpoints as unavailable
start_time = time.time()
while (len(setup[COLLECTION].client_connection._global_endpoint_manager.location_cache.location_unavailability_info_by_endpoint) < len(
REGIONS) and time.time() - start_time < 10):
await asyncio.sleep(0.1)
finally:
_global_endpoint_manager_async._GlobalEndpointManager._GetDatabaseAccountStub = self.original_getDatabaseAccountStub
setup[COLLECTION].client_connection.connection_policy.PreferredLocations = self.original_preferred_locations
_cosmos_client_connection_async.CosmosClientConnection._GetDatabaseAccountCheck = self.original_getDatabaseAccountCheck
_retry_utility_async.ExecuteFunctionAsync = self.OriginalExecuteFunction
num_unavailable_endpoints = len(REGIONS)
unavailable_endpoint_info = setup[COLLECTION].client_connection._global_endpoint_manager.location_cache.location_unavailability_info_by_endpoint
assert len(unavailable_endpoint_info) == num_unavailable_endpoints
write_endpoints = {
_location_cache.LocationCache.GetLocationalEndpoint(self.host, REGION_1)
}
for endpoint, info in unavailable_endpoint_info.items():
assert _location_cache.EndpointOperationType.ReadType in info["operationType"]
if endpoint in write_endpoints:
assert _location_cache.EndpointOperationType.WriteType in info["operationType"]
else:
assert _location_cache.EndpointOperationType.WriteType not in info["operationType"]
async def mock_health_check(self, **kwargs):
await asyncio.sleep(100)
async def mock_health_check_failure(self, **kwargs):
await asyncio.sleep(1)
raise Exception("Mock health check failure")
class MockGetDatabaseAccountCheck(object):
def __init__(self):
self.counter = 0
self.index = 0
async def __call__(self, endpoint):
self.index += 1
self.counter += 1
class MockGetDatabaseAccountCheckError(object):
async def __call__(self, endpoint):
raise ServiceRequestError("Mock health check failure")
class MockGetDatabaseAccount(object):
def __init__(
self,
regions: List[str],
):
self.regions = regions
async def __call__(self, endpoint):
read_regions = self.regions
read_locations = []
for loc in read_regions:
locational_endpoint = _location_cache.LocationCache.GetLocationalEndpoint(endpoint, loc)
account_endpoint = locational_endpoint
read_locations.append({'databaseAccountEndpoint': account_endpoint, 'name': loc})
write_regions = [self.regions[0]]
write_locations = []
for loc in write_regions:
locational_endpoint = _location_cache.LocationCache.GetLocationalEndpoint(endpoint, loc)
account_endpoint = locational_endpoint
write_locations.append({'databaseAccountEndpoint': account_endpoint, 'name': loc})
multi_write = False
db_acc = DatabaseAccount()
db_acc.DatabasesLink = "/dbs/"
db_acc.MediaLink = "/media/"
db_acc._ReadableLocations = read_locations
db_acc._WritableLocations = write_locations
db_acc._EnableMultipleWritableLocations = multi_write
db_acc.ConsistencyPolicy = {"defaultConsistencyLevel": "Session"}
return db_acc
if __name__ == '__main__':
unittest.main()
|
