Skip to content

Commit ed5a541

Browse files
tvaron3kushagraThapar
tvaron3
and
kushagraThapar
authored
Change how util functions imported (#41449)
* fix release? * revert change * fix tests * Fixed imports --------- Co-authored-by: Kushagra Thapar <[email protected]>
1 parent cb2790e commit ed5a541

6 files changed

+151
-154
lines changed

‎sdk/cosmos/azure-cosmos/tests/test_circuit_breaker_emulator.py‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -13,11 +13,10 @@
1313
from azure.cosmos.exceptions import CosmosHttpResponseError
1414
from _fault_injection_transport import FaultInjectionTransport
1515
from azure.cosmos.http_constants import ResourceType
16-
from test_per_partition_circuit_breaker_mm import perform_write_operation
17-
from test_per_partition_circuit_breaker_mm_async import (create_doc, PK_VALUE, create_errors,
18-
DELETE_ALL_ITEMS_BY_PARTITION_KEY,
19-
validate_unhealthy_partitions as validate_unhealthy_partitions_mm)
20-
from test_per_partition_circuit_breaker_sm_mrr_async import validate_unhealthy_partitions as validate_unhealthy_partitions_sm_mrr
16+
from test_per_partition_circuit_breaker_mm import create_doc, DELETE_ALL_ITEMS_BY_PARTITION_KEY, PK_VALUE, \
17+
create_errors, perform_write_operation, validate_unhealthy_partitions as validate_unhealthy_partitions_mm
18+
from test_per_partition_circuit_breaker_sm_mrr import \
19+
validate_unhealthy_partitions as validate_unhealthy_partitions_sm_mrr
2120

2221
COLLECTION = "created_collection"
2322
@pytest.fixture(scope="class", autouse=True)

‎sdk/cosmos/azure-cosmos/tests/test_circuit_breaker_emulator_async.py‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -14,11 +14,11 @@
1414
from azure.cosmos.aio import CosmosClient
1515
from azure.cosmos.exceptions import CosmosHttpResponseError
1616
from azure.cosmos.http_constants import ResourceType
17-
from test_per_partition_circuit_breaker_mm_async import (create_doc, PK_VALUE, create_errors,
18-
DELETE_ALL_ITEMS_BY_PARTITION_KEY,
19-
validate_unhealthy_partitions as validate_unhealthy_partitions_mm,
20-
perform_write_operation, cleanup_method)
21-
from test_per_partition_circuit_breaker_sm_mrr_async import validate_unhealthy_partitions as validate_unhealthy_partitions_sm_mrr
17+
from test_per_partition_circuit_breaker_mm_async import (perform_write_operation, cleanup_method)
18+
from test_per_partition_circuit_breaker_mm import create_doc, DELETE_ALL_ITEMS_BY_PARTITION_KEY, PK_VALUE, \
19+
create_errors, validate_unhealthy_partitions as validate_unhealthy_partitions_mm
20+
from test_per_partition_circuit_breaker_sm_mrr import \
21+
validate_unhealthy_partitions as validate_unhealthy_partitions_sm_mrr
2222
from _fault_injection_transport_async import FaultInjectionTransportAsync
2323

2424
COLLECTION = "created_collection"

‎sdk/cosmos/azure-cosmos/tests/test_per_partition_circuit_breaker_mm.py‎

Lines changed: 101 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -11,15 +11,92 @@
1111
import test_config
1212
from azure.cosmos import _location_cache, _partition_health_tracker
1313
from azure.cosmos import CosmosClient
14+
from azure.cosmos._partition_health_tracker import HEALTH_STATUS, UNHEALTHY_TENTATIVE, UNHEALTHY
1415
from azure.cosmos.exceptions import CosmosHttpResponseError
1516
from _fault_injection_transport import FaultInjectionTransport
16-
from test_per_partition_circuit_breaker_mm_async import DELETE, CREATE, UPSERT, REPLACE, PATCH, BATCH, \
17-
validate_response_uri, READ, \
18-
QUERY_PK, QUERY, CHANGE_FEED, CHANGE_FEED_PK, CHANGE_FEED_EPK, READ_ALL_ITEMS, REGION_1, REGION_2, \
19-
write_operations_and_errors, validate_unhealthy_partitions, read_operations_and_errors, PK_VALUE, operations, \
20-
create_doc, validate_stats
21-
from test_per_partition_circuit_breaker_mm_async import DELETE_ALL_ITEMS_BY_PARTITION_KEY
2217

18+
REGION_1 = "West US 3"
19+
REGION_2 = "West US"
20+
CHANGE_FEED = "changefeed"
21+
CHANGE_FEED_PK = "changefeed_pk"
22+
CHANGE_FEED_EPK = "changefeed_epk"
23+
READ = "read"
24+
CREATE = "create"
25+
READ_ALL_ITEMS = "read_all_items"
26+
DELETE_ALL_ITEMS_BY_PARTITION_KEY = "delete_all_items_by_partition_key"
27+
QUERY = "query"
28+
QUERY_PK = "query_pk"
29+
BATCH = "batch"
30+
UPSERT = "upsert"
31+
REPLACE = "replace"
32+
PATCH = "patch"
33+
DELETE = "delete"
34+
PK_VALUE = "pk1"
35+
36+
def create_doc():
37+
return {'id': str(uuid.uuid4()),
38+
'pk': PK_VALUE,
39+
'name': 'sample document',
40+
'key': 'value'}
41+
42+
43+
def read_operations_and_errors():
44+
read_operations = [READ, QUERY_PK, CHANGE_FEED, CHANGE_FEED_PK, CHANGE_FEED_EPK]
45+
errors = create_errors()
46+
params = []
47+
for read_operation in read_operations:
48+
for error in errors:
49+
params.append((read_operation, error))
50+
51+
return params
52+
53+
def write_operations_and_errors():
54+
write_operations = [CREATE, UPSERT, REPLACE, DELETE, PATCH, BATCH]
55+
errors = create_errors()
56+
params = []
57+
for write_operation in write_operations:
58+
for error in errors:
59+
params.append((write_operation, error))
60+
61+
return params
62+
63+
def operations():
64+
write_operations = [CREATE, UPSERT, REPLACE, DELETE, PATCH, BATCH]
65+
read_operations = [READ, QUERY_PK, CHANGE_FEED_PK, CHANGE_FEED_EPK]
66+
operations = []
67+
for i, write_operation in enumerate(write_operations):
68+
operations.append((read_operations[i % len(read_operations)], write_operation))
69+
70+
return operations
71+
72+
def create_errors():
73+
errors = []
74+
error_codes = [408, 500, 502, 503]
75+
for error_code in error_codes:
76+
errors.append(CosmosHttpResponseError(
77+
status_code=error_code,
78+
message="Some injected error."))
79+
errors.append(ServiceResponseError(message="Injected Service Response Error."))
80+
return errors
81+
82+
def validate_unhealthy_partitions(global_endpoint_manager,
83+
expected_unhealthy_partitions):
84+
health_info_map = global_endpoint_manager.global_partition_endpoint_manager_core.partition_health_tracker.pk_range_wrapper_to_health_info
85+
unhealthy_partitions = 0
86+
for pk_range_wrapper, location_to_health_info in health_info_map.items():
87+
for location, health_info in location_to_health_info.items():
88+
health_status = health_info.unavailability_info.get(HEALTH_STATUS)
89+
if health_status == UNHEALTHY_TENTATIVE or health_status == UNHEALTHY:
90+
unhealthy_partitions += 1
91+
else:
92+
assert health_info.read_consecutive_failure_count < 10
93+
assert health_info.write_consecutive_failure_count < 5
94+
95+
assert unhealthy_partitions == expected_unhealthy_partitions
96+
97+
def validate_response_uri(response, expected_uri):
98+
request = response.get_response_headers()["_request"]
99+
assert request.url.startswith(expected_uri)
23100
def perform_write_operation(operation, container, fault_injection_container, doc_id, pk, expected_uri):
24101
doc = {'id': doc_id,
25102
'pk': pk,
@@ -414,3 +491,21 @@ def test_service_request_error(self, read_operation, write_operation):
414491

415492
if __name__ == '__main__':
416493
unittest.main()
494+
495+
496+
def validate_stats(global_endpoint_manager,
497+
expected_write_consecutive_failure_count,
498+
expected_read_consecutive_failure_count,
499+
expected_read_failure_count,
500+
expected_write_failure_count,
501+
expected_write_success_count,
502+
expected_read_success_count):
503+
health_info_map = global_endpoint_manager.global_partition_endpoint_manager_core.partition_health_tracker.pk_range_wrapper_to_health_info
504+
for pk_range_wrapper, location_to_health_info in health_info_map.items():
505+
health_info = location_to_health_info[REGION_1]
506+
assert health_info.read_consecutive_failure_count == expected_read_consecutive_failure_count
507+
assert health_info.write_consecutive_failure_count == expected_write_consecutive_failure_count
508+
assert health_info.read_failure_count == expected_read_failure_count
509+
assert health_info.write_failure_count == expected_write_failure_count
510+
assert health_info.read_success_count == expected_read_success_count
511+
assert health_info.write_success_count == expected_write_success_count

‎sdk/cosmos/azure-cosmos/tests/test_per_partition_circuit_breaker_mm_async.py‎

Lines changed: 14 additions & 110 deletions
Original file line numberDiff line numberDiff line change
@@ -12,81 +12,17 @@
1212

1313
import test_config
1414
from azure.cosmos import _location_cache, _partition_health_tracker
15-
from azure.cosmos._partition_health_tracker import HEALTH_STATUS, UNHEALTHY, UNHEALTHY_TENTATIVE
1615
from azure.cosmos.aio import CosmosClient
1716
from azure.cosmos.exceptions import CosmosHttpResponseError
1817
from _fault_injection_transport_async import FaultInjectionTransportAsync
19-
20-
REGION_1 = "West US 3"
21-
REGION_2 = "West US"
22-
CHANGE_FEED = "changefeed"
23-
CHANGE_FEED_PK = "changefeed_pk"
24-
CHANGE_FEED_EPK = "changefeed_epk"
25-
READ = "read"
26-
CREATE = "create"
27-
READ_ALL_ITEMS = "read_all_items"
28-
DELETE_ALL_ITEMS_BY_PARTITION_KEY = "delete_all_items_by_partition_key"
29-
QUERY = "query"
30-
QUERY_PK = "query_pk"
31-
BATCH = "batch"
32-
UPSERT = "upsert"
33-
REPLACE = "replace"
34-
PATCH = "patch"
35-
DELETE = "delete"
36-
PK_VALUE = "pk1"
37-
18+
from test_per_partition_circuit_breaker_mm import create_doc, read_operations_and_errors, \
19+
write_operations_and_errors, operations, REGION_1, REGION_2, CHANGE_FEED, CHANGE_FEED_PK, CHANGE_FEED_EPK, READ, \
20+
CREATE, READ_ALL_ITEMS, DELETE_ALL_ITEMS_BY_PARTITION_KEY, QUERY, QUERY_PK, BATCH, UPSERT, REPLACE, PATCH, DELETE, \
21+
PK_VALUE, validate_unhealthy_partitions, validate_response_uri
22+
from test_per_partition_circuit_breaker_mm import validate_stats
3823

3924
COLLECTION = "created_collection"
4025

41-
def create_errors():
42-
errors = []
43-
error_codes = [408, 500, 502, 503]
44-
for error_code in error_codes:
45-
errors.append(CosmosHttpResponseError(
46-
status_code=error_code,
47-
message="Some injected error."))
48-
errors.append(ServiceResponseError(message="Injected Service Response Error."))
49-
return errors
50-
51-
def write_operations_and_errors():
52-
write_operations = [CREATE, UPSERT, REPLACE, DELETE, PATCH, BATCH]
53-
errors = create_errors()
54-
params = []
55-
for write_operation in write_operations:
56-
for error in errors:
57-
params.append((write_operation, error))
58-
59-
return params
60-
61-
def create_doc():
62-
return {'id': str(uuid.uuid4()),
63-
'pk': PK_VALUE,
64-
'name': 'sample document',
65-
'key': 'value'}
66-
67-
def read_operations_and_errors():
68-
read_operations = [READ, QUERY_PK, CHANGE_FEED, CHANGE_FEED_PK, CHANGE_FEED_EPK]
69-
errors = create_errors()
70-
params = []
71-
for read_operation in read_operations:
72-
for error in errors:
73-
params.append((read_operation, error))
74-
75-
return params
76-
77-
def operations():
78-
write_operations = [CREATE, UPSERT, REPLACE, DELETE, PATCH, BATCH]
79-
read_operations = [READ, QUERY_PK, CHANGE_FEED_PK, CHANGE_FEED_EPK]
80-
operations = []
81-
for i, write_operation in enumerate(write_operations):
82-
operations.append((read_operations[i % len(read_operations)], write_operation))
83-
84-
return operations
85-
86-
def validate_response_uri(response, expected_uri):
87-
request = response.get_response_headers()["_request"]
88-
assert request.url.startswith(expected_uri)
89-
9026
async def perform_write_operation(operation, container, fault_injection_container, doc_id, pk, expected_uri):
9127
doc = {'id': doc_id,
9228
'pk': pk,
@@ -162,38 +98,6 @@ async def perform_read_operation(operation, container, doc_id, pk, expected_uri)
16298
async for _ in container.read_all_items():
16399
pass
164100

165-
def validate_unhealthy_partitions(global_endpoint_manager,
166-
expected_unhealthy_partitions):
167-
health_info_map = global_endpoint_manager.global_partition_endpoint_manager_core.partition_health_tracker.pk_range_wrapper_to_health_info
168-
unhealthy_partitions = 0
169-
for pk_range_wrapper, location_to_health_info in health_info_map.items():
170-
for location, health_info in location_to_health_info.items():
171-
health_status = health_info.unavailability_info.get(HEALTH_STATUS)
172-
if health_status == UNHEALTHY_TENTATIVE or health_status == UNHEALTHY:
173-
unhealthy_partitions += 1
174-
else:
175-
assert health_info.read_consecutive_failure_count < 10
176-
assert health_info.write_consecutive_failure_count < 5
177-
178-
assert unhealthy_partitions == expected_unhealthy_partitions
179-
180-
def validate_stats(global_endpoint_manager,
181-
expected_write_consecutive_failure_count,
182-
expected_read_consecutive_failure_count,
183-
expected_read_failure_count,
184-
expected_write_failure_count,
185-
expected_write_success_count,
186-
expected_read_success_count):
187-
health_info_map = global_endpoint_manager.global_partition_endpoint_manager_core.partition_health_tracker.pk_range_wrapper_to_health_info
188-
for pk_range_wrapper, location_to_health_info in health_info_map.items():
189-
health_info = location_to_health_info[REGION_1]
190-
assert health_info.read_consecutive_failure_count == expected_read_consecutive_failure_count
191-
assert health_info.write_consecutive_failure_count == expected_write_consecutive_failure_count
192-
assert health_info.read_failure_count == expected_read_failure_count
193-
assert health_info.write_failure_count == expected_write_failure_count
194-
assert health_info.read_success_count == expected_read_success_count
195-
assert health_info.write_success_count == expected_write_success_count
196-
197101
async def cleanup_method(initialized_objects: List[Dict[str, Any]]):
198102
for obj in initialized_objects:
199103
method_client: CosmosClient = obj["client"]
@@ -416,10 +320,10 @@ async def test_read_failure_rate_threshold_async(self, read_operation, error):
416320
validate_unhealthy_partitions(global_endpoint_manager, 0)
417321
# read will fail and retry in other region
418322
await perform_read_operation(read_operation,
419-
fault_injection_container,
420-
doc['id'],
421-
PK_VALUE,
422-
expected_uri)
323+
fault_injection_container,
324+
doc['id'],
325+
PK_VALUE,
326+
expected_uri)
423327
if read_operation in (CHANGE_FEED, QUERY, READ_ALL_ITEMS):
424328
# these operations are cross partition so they would mark both partitions as unavailable
425329
expected_unhealthy_partitions = 5
@@ -468,7 +372,7 @@ async def test_stat_reset_async(self):
468372
except CosmosHttpResponseError as e:
469373
assert e.status_code == 503
470374
validate_unhealthy_partitions(global_endpoint_manager, 0)
471-
validate_stats(global_endpoint_manager, 2, 2, 2, 2, 0, 0)
375+
validate_stats(global_endpoint_manager, 2, 2, 2, 2, 0, 0)
472376
await asyncio.sleep(25)
473377
await perform_read_operation(READ,
474378
fault_injection_container,
@@ -507,10 +411,10 @@ async def test_service_request_error_async(self, read_operation, write_operation
507411
custom_transport.faults = []
508412
try:
509413
await perform_read_operation(read_operation,
510-
fault_injection_container,
511-
doc['id'],
512-
PK_VALUE,
513-
expected_uri)
414+
fault_injection_container,
415+
doc['id'],
416+
PK_VALUE,
417+
expected_uri)
514418
finally:
515419
_partition_health_tracker.INITIAL_UNAVAILABLE_TIME_MS = original_unavailable_time
516420
validate_unhealthy_partitions(global_endpoint_manager, 0)

‎sdk/cosmos/azure-cosmos/tests/test_per_partition_circuit_breaker_sm_mrr.py‎

Lines changed: 23 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -11,15 +11,32 @@
1111
import test_config
1212
from azure.cosmos import _partition_health_tracker, _location_cache
1313
from azure.cosmos import CosmosClient
14+
from azure.cosmos._partition_health_tracker import HEALTH_STATUS, UNHEALTHY_TENTATIVE, UNHEALTHY
1415
from azure.cosmos.exceptions import CosmosHttpResponseError
1516
from _fault_injection_transport import FaultInjectionTransport
16-
from test_per_partition_circuit_breaker_mm import perform_write_operation, perform_read_operation
17-
from test_per_partition_circuit_breaker_mm_async import create_doc, PK_VALUE, write_operations_and_errors, \
18-
operations, REGION_2, REGION_1, READ, CREATE, validate_stats
19-
from test_per_partition_circuit_breaker_sm_mrr_async import validate_unhealthy_partitions
17+
from test_per_partition_circuit_breaker_mm import create_doc, write_operations_and_errors, operations, REGION_1, \
18+
REGION_2, PK_VALUE, perform_write_operation, perform_read_operation
2019

2120
COLLECTION = "created_collection"
2221

22+
def validate_unhealthy_partitions(global_endpoint_manager,
23+
expected_unhealthy_partitions):
24+
health_info_map = global_endpoint_manager.global_partition_endpoint_manager_core.partition_health_tracker.pk_range_wrapper_to_health_info
25+
unhealthy_partitions = 0
26+
for pk_range_wrapper, location_to_health_info in health_info_map.items():
27+
for location, health_info in location_to_health_info.items():
28+
health_status = health_info.unavailability_info.get(HEALTH_STATUS)
29+
if health_status == UNHEALTHY_TENTATIVE or health_status == UNHEALTHY:
30+
unhealthy_partitions += 1
31+
32+
else:
33+
assert health_info.read_consecutive_failure_count < 10
34+
# single region write account should never track write failures
35+
assert health_info.write_failure_count == 0
36+
assert health_info.write_consecutive_failure_count == 0
37+
38+
assert unhealthy_partitions == expected_unhealthy_partitions
39+
2340
@pytest.mark.cosmosCircuitBreakerMultiRegion
2441
class TestPerPartitionCircuitBreakerSmMrr:
2542
host = test_config.TestConfig.host
@@ -219,4 +236,5 @@ def test_service_request_error(self, read_operation, write_operation):
219236
# test cosmos client timeout
220237

221238
if __name__ == '__main__':
222-
unittest.main()
239+
unittest.main()
240+

0 commit comments

Comments
 (0)