From d75299b288e2f89d3989f9036e3d13c6ea769995 Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Wed, 29 Apr 2026 16:46:01 -0400 Subject: [PATCH 01/22] Update test_per_partition_automatic_failover_async.py --- .../tests/test_per_partition_automatic_failover_async.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/tests/test_per_partition_automatic_failover_async.py b/sdk/cosmos/azure-cosmos/tests/test_per_partition_automatic_failover_async.py index 860727c18b0b..5bc268ca8607 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_per_partition_automatic_failover_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_per_partition_automatic_failover_async.py @@ -61,7 +61,7 @@ async def setup_info(self, error=None, max_count=None, is_batch=False, exclude_c doc_fail_id = str(uuid.uuid4()) doc_success_id = str(uuid.uuid4()) predicate = lambda r: (FaultInjectionTransportAsync.predicate_req_for_document_with_id(r, doc_fail_id) and - FaultInjectionTransportAsync.predicate_is_write_operation(r, "com")) + FaultInjectionTransportAsync.predicate_is_write_operation(r, "documents")) # The MockRequest only gets used to create the MockHttpResponse mock_request = FaultInjectionTransport.MockHttpRequest(url=self.host) if is_batch: From b7b0cba84f687c604fe12eaf29f14b8aacbdc6f0 Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Wed, 29 Apr 2026 17:25:10 -0400 Subject: [PATCH 02/22] ignore semantic reranker tests for staging --- sdk/cosmos/azure-cosmos/pytest.ini | 1 + sdk/cosmos/cosmos-staging.yml | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/pytest.ini b/sdk/cosmos/azure-cosmos/pytest.ini index a5e006ea027e..689522387d65 100644 --- a/sdk/cosmos/azure-cosmos/pytest.ini +++ b/sdk/cosmos/azure-cosmos/pytest.ini @@ -8,3 +8,4 @@ markers = cosmosCircuitBreaker: marks tests running on Cosmos DB live account with per partition circuit breaker enabled and multi-write enabled. cosmosCircuitBreakerMultiRegion: marks tests running on Cosmos DB live account with one write region and multiple read regions and per partition circuit breaker enabled. cosmosPerPartitionAutomaticFailover: marks tests running on Cosmos DB live account with one write region and multiple read regions and per partition automatic failover enabled. + semanticReranker: marks tests running on a Cosmos DB live account with semantic reranker enabled. diff --git a/sdk/cosmos/cosmos-staging.yml b/sdk/cosmos/cosmos-staging.yml index c8f098a33125..333a1536a69d 100644 --- a/sdk/cosmos/cosmos-staging.yml +++ b/sdk/cosmos/cosmos-staging.yml @@ -37,7 +37,7 @@ jobs: script: >- cd ./sdk/cosmos/azure-cosmos - azpysdk whl --pytest-args -k 'not cosmosEmulator and not cosmosMultiRegion and not cosmosCircuitBreaker and not cosmosCircuitBreakerMultiRegion' + azpysdk whl --pytest-args -k 'not cosmosEmulator and not cosmosMultiRegion and not cosmosCircuitBreaker and not cosmosCircuitBreakerMultiRegion and not semanticReranker' env: ACCOUNT_KEY: $(ACCOUNT_KEY) # Maps the secret variable $(token) from my-var-group - task: AzurePowerShell@5 From 429f935eb291f770dbefe71b8bb859ea46e290a3 Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Wed, 29 Apr 2026 19:29:21 -0400 Subject: [PATCH 03/22] Update test_per_partition_automatic_failover.py --- .../azure-cosmos/tests/test_per_partition_automatic_failover.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/tests/test_per_partition_automatic_failover.py b/sdk/cosmos/azure-cosmos/tests/test_per_partition_automatic_failover.py index 437c25556e05..30dbdc6bee42 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_per_partition_automatic_failover.py +++ b/sdk/cosmos/azure-cosmos/tests/test_per_partition_automatic_failover.py @@ -66,7 +66,7 @@ def setup_info(self, error=None, max_count=None, is_batch=False, exclude_client_ doc_fail_id = str(uuid.uuid4()) doc_success_id = str(uuid.uuid4()) predicate = lambda r: (FaultInjectionTransport.predicate_req_for_document_with_id(r, doc_fail_id) and - FaultInjectionTransport.predicate_is_write_operation(r, "com")) + FaultInjectionTransport.predicate_is_write_operation(r, "documents")) # The MockRequest only gets used to create the MockHttpResponse mock_request = FaultInjectionTransport.MockHttpRequest(url=self.host) if is_batch: From 84cc193cd771ff9e5a8aae8a7087f0edcc4da5fe Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Thu, 30 Apr 2026 12:38:53 -0400 Subject: [PATCH 04/22] rename and fix existing, add new container RID persistence tests --- .../tests/test_container_rid_header_unit.py | 49 +++++++++++++++- .../test_container_rid_header_unit_async.py | 57 ++++++++++++++++++- 2 files changed, 100 insertions(+), 6 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/tests/test_container_rid_header_unit.py b/sdk/cosmos/azure-cosmos/tests/test_container_rid_header_unit.py index d9dcbd7068c1..8b155a44fbcf 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_container_rid_header_unit.py +++ b/sdk/cosmos/azure-cosmos/tests/test_container_rid_header_unit.py @@ -149,7 +149,7 @@ def test_end_to_end_without_containerRID_omits_header(self): ) assert "containerRID" not in client.captured_feed_options - def test_force_refresh_preserves_containerRID(self): + def test_incremental_refresh_preserves_containerRID(self): """A force_refresh (triggered by split/merge errors) must re-send containerRID to the service so the correct physical container is queried again.""" @@ -165,8 +165,51 @@ def test_force_refresh_preserves_containerRID(self): force_refresh=True, previous_routing_map=previous_map, ) - # In this mock setup, incremental refresh has no split/merge deltas and - # falls back to a full refresh, so force-refresh performs two reads. + # The mock returns the same ranges (no split/merge deltas), so the + # incremental merge succeeds in a single fetch without falling back + # to a full refresh. + assert client.call_count == 2 + assert client.captured_feed_options.get("containerRID") == CONTAINER_RID + + def test_full_refresh_fallback_after_split_preserves_containerRID(self): + """When an incremental refresh encounters unresolvable parent ranges, + it falls back to a full refresh. containerRID must still flow through + on the fallback call.""" + # Post-split ranges reference parents not in the original map + split_ranges = [ + {"id": "3", "minInclusive": "", "maxExclusive": "1F", "parents": ["99"]}, + {"id": "4", "minInclusive": "1F", "maxExclusive": "3F", "parents": ["99"]}, + ] + + class SplitMockClient: + def __init__(self): + self.captured_feed_options = None + self.call_count = 0 + + def _ReadPartitionKeyRanges(self, _collection_link, feed_options=None, **kwargs): + self.captured_feed_options = dict(feed_options) if feed_options else {} + self.call_count += 1 + # First call: initial load returns original ranges + # Second call: incremental returns split ranges with unknown parents + # Third call: full fallback returns original ranges + if self.call_count == 2: + return iter(split_ranges) + return iter(PARTITION_KEY_RANGES) + + client = SplitMockClient() + cache = PartitionKeyRangeCache(client) + feed_options = {"containerRID": CONTAINER_RID} + previous_map = cache.get_routing_map(COLLECTION_LINK, feed_options) + assert client.call_count == 1 + assert client.captured_feed_options.get("containerRID") == CONTAINER_RID + cache.get_routing_map( + COLLECTION_LINK, + feed_options, + force_refresh=True, + previous_routing_map=previous_map, + ) + # Incremental fetch returns ranges with unresolvable parents, triggering + # a full refresh fallback: 1 initial + 1 incremental + 1 full = 3 calls. assert client.call_count == 3 assert client.captured_feed_options.get("containerRID") == CONTAINER_RID diff --git a/sdk/cosmos/azure-cosmos/tests/test_container_rid_header_unit_async.py b/sdk/cosmos/azure-cosmos/tests/test_container_rid_header_unit_async.py index 698328982a80..5de1cfbb9379 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_container_rid_header_unit_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_container_rid_header_unit_async.py @@ -139,7 +139,7 @@ async def test_end_to_end_without_containerRID_omits_header_async(self): ) assert "containerRID" not in client.captured_feed_options - async def test_force_refresh_preserves_containerRID_async(self): + async def test_incremental_refresh_preserves_containerRID_async(self): """A force_refresh (triggered by split/merge errors) must re-send containerRID to the service so the correct physical container is queried again.""" @@ -155,8 +155,59 @@ async def test_force_refresh_preserves_containerRID_async(self): force_refresh=True, previous_routing_map=previous_map, ) - # In this mock setup, incremental refresh has no split/merge deltas and - # falls back to a full refresh, so force-refresh performs two reads. + # The mock returns the same ranges (no split/merge deltas), so the + # incremental merge succeeds in a single fetch without falling back + # to a full refresh. + assert client.call_count == 2 + assert client.captured_feed_options.get("containerRID") == CONTAINER_RID + + async def test_full_refresh_fallback_after_split_preserves_containerRID_async(self): + """When an incremental refresh encounters unresolvable parent ranges, + it falls back to a full refresh. containerRID must still flow through + on the fallback call.""" + # Post-split ranges reference parents not in the original map + split_ranges = [ + {"id": "3", "minInclusive": "", "maxExclusive": "1F", "parents": ["99"]}, + {"id": "4", "minInclusive": "1F", "maxExclusive": "3F", "parents": ["99"]}, + ] + + class SplitMockClient: + def __init__(self): + self.captured_feed_options = None + self.call_count = 0 + + @staticmethod + def _capture_internal_headers(kwargs, etag): + captured_headers = kwargs.get('_internal_response_headers_capture') + if captured_headers is not None: + captured_headers.clear() + captured_headers.update({'ETag': etag}) + + async def _ReadPartitionKeyRanges(self, _collection_link, feed_options=None, **kwargs): + self.captured_feed_options = dict(feed_options) if feed_options else {} + self.call_count += 1 + SplitMockClient._capture_internal_headers(kwargs, '"test-etag-1"') + # First call: initial load returns original ranges + # Second call: incremental returns split ranges with unknown parents + # Third call: full fallback returns original ranges + items = split_ranges if self.call_count == 2 else PARTITION_KEY_RANGES + for item in items: + yield item + + client = SplitMockClient() + cache = PartitionKeyRangeCache(client) + feed_options: Dict[str, Any] = {"containerRID": CONTAINER_RID} + previous_map = await cache.get_routing_map(COLLECTION_LINK, feed_options) + assert client.call_count == 1 + assert client.captured_feed_options.get("containerRID") == CONTAINER_RID + await cache.get_routing_map( + COLLECTION_LINK, + feed_options, + force_refresh=True, + previous_routing_map=previous_map, + ) + # Incremental fetch returns ranges with unresolvable parents, triggering + # a full refresh fallback: 1 initial + 1 incremental + 1 full = 3 calls. assert client.call_count == 3 assert client.captured_feed_options.get("containerRID") == CONTAINER_RID From 6ea07255b43f9d74e2c0b3435e76b4c20d6237b9 Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Thu, 30 Apr 2026 15:03:10 -0400 Subject: [PATCH 05/22] Update cspell.json --- sdk/cosmos/azure-cosmos/cspell.json | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/cspell.json b/sdk/cosmos/azure-cosmos/cspell.json index 5eb68afe9de1..1caf92d0589e 100644 --- a/sdk/cosmos/azure-cosmos/cspell.json +++ b/sdk/cosmos/azure-cosmos/cspell.json @@ -4,6 +4,8 @@ "perfresults", "ppcb", "hdrh", - "hdrhistogram" + "hdrhistogram", + "reranker", + "Reranker" ] } From 8c908877ed809416ed47d20534110e6bb3953cda Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Thu, 30 Apr 2026 15:07:41 -0400 Subject: [PATCH 06/22] again service updates to error messages -_- --- sdk/cosmos/azure-cosmos/tests/test_full_text_policy.py | 2 +- sdk/cosmos/azure-cosmos/tests/test_full_text_policy_async.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/tests/test_full_text_policy.py b/sdk/cosmos/azure-cosmos/tests/test_full_text_policy.py index 2f739fcb5b34..64e5f587a1af 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_full_text_policy.py +++ b/sdk/cosmos/azure-cosmos/tests/test_full_text_policy.py @@ -231,7 +231,7 @@ def test_fail_create_full_text_policy(self): except exceptions.CosmosHttpResponseError as e: assert e.status_code == 400 assert re.search( - r"the full.text policy contains an unsupported language.*spa-SPA", + r"The language specified in the full-text policy, \'spa-SPA\', is invalid.", e.http_error_message, re.IGNORECASE) # Pass a full text policy with an unsupported path language diff --git a/sdk/cosmos/azure-cosmos/tests/test_full_text_policy_async.py b/sdk/cosmos/azure-cosmos/tests/test_full_text_policy_async.py index c962c5d398e8..cb34844ced9b 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_full_text_policy_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_full_text_policy_async.py @@ -245,7 +245,7 @@ async def test_fail_create_full_text_policy_async(self): except exceptions.CosmosHttpResponseError as e: assert e.status_code == 400 assert re.search( - r"the full.text policy contains an unsupported language.*spa-SPA", + r"The language specified in the full-text policy, \'spa-SPA\', is invalid.", e.http_error_message, re.IGNORECASE) # Pass a full text policy with an unsupported path language From 25a6d080be6304ed259b64225d9d8395b51ff3b4 Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Thu, 30 Apr 2026 16:08:30 -0400 Subject: [PATCH 07/22] cspell --- sdk/cosmos/azure-cosmos/cspell.json | 3 +-- sdk/cosmos/cosmos-staging.yml | 2 +- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/cspell.json b/sdk/cosmos/azure-cosmos/cspell.json index 1caf92d0589e..5403accddff6 100644 --- a/sdk/cosmos/azure-cosmos/cspell.json +++ b/sdk/cosmos/azure-cosmos/cspell.json @@ -5,7 +5,6 @@ "ppcb", "hdrh", "hdrhistogram", - "reranker", - "Reranker" + "reranker" ] } diff --git a/sdk/cosmos/cosmos-staging.yml b/sdk/cosmos/cosmos-staging.yml index 333a1536a69d..5cb9e090d290 100644 --- a/sdk/cosmos/cosmos-staging.yml +++ b/sdk/cosmos/cosmos-staging.yml @@ -37,7 +37,7 @@ jobs: script: >- cd ./sdk/cosmos/azure-cosmos - azpysdk whl --pytest-args -k 'not cosmosEmulator and not cosmosMultiRegion and not cosmosCircuitBreaker and not cosmosCircuitBreakerMultiRegion and not semanticReranker' + azpysdk whl --pytest-args -k 'not cosmosEmulator and not cosmosMultiRegion and not cosmosCircuitBreaker and not cosmosCircuitBreakerMultiRegion and not semanticReranker' # cspell:ignore Reranker env: ACCOUNT_KEY: $(ACCOUNT_KEY) # Maps the secret variable $(token) from my-var-group - task: AzurePowerShell@5 From 51babcabac37ca88c78150d369064f7d5a9f3583 Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Thu, 30 Apr 2026 16:12:51 -0400 Subject: [PATCH 08/22] Update cspell.json --- sdk/cosmos/azure-cosmos/cspell.json | 3 --- 1 file changed, 3 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/cspell.json b/sdk/cosmos/azure-cosmos/cspell.json index 5e6b7ecbb334..fd2c6923e829 100644 --- a/sdk/cosmos/azure-cosmos/cspell.json +++ b/sdk/cosmos/azure-cosmos/cspell.json @@ -6,10 +6,7 @@ "perfresults", "pkrange", "ppcb", - "hdrh", - "hdrhistogram", "reranker", - "Reranker" "toctou" ] } From 720afd0265e7ebc43f31dcc9691468d27717b42d Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Thu, 30 Apr 2026 18:06:46 -0400 Subject: [PATCH 09/22] full text policies again --- sdk/cosmos/azure-cosmos/tests/test_full_text_policy.py | 2 +- sdk/cosmos/azure-cosmos/tests/test_full_text_policy_async.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/tests/test_full_text_policy.py b/sdk/cosmos/azure-cosmos/tests/test_full_text_policy.py index 64e5f587a1af..c70906844aca 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_full_text_policy.py +++ b/sdk/cosmos/azure-cosmos/tests/test_full_text_policy.py @@ -254,7 +254,7 @@ def test_fail_create_full_text_policy(self): except exceptions.CosmosHttpResponseError as e: assert e.status_code == 400 assert re.search( - r"the full.text policy contains an unsupported language.*spa-SPA", + r"The language specified in the full-text policy, \'spa-SPA\', is invalid", e.http_error_message, re.IGNORECASE) def test_fail_create_full_text_indexing_policy(self): diff --git a/sdk/cosmos/azure-cosmos/tests/test_full_text_policy_async.py b/sdk/cosmos/azure-cosmos/tests/test_full_text_policy_async.py index cb34844ced9b..c4b67f07f2ef 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_full_text_policy_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_full_text_policy_async.py @@ -268,7 +268,7 @@ async def test_fail_create_full_text_policy_async(self): except exceptions.CosmosHttpResponseError as e: assert e.status_code == 400 assert re.search( - r"the full.text policy contains an unsupported language.*spa-SPA", + r"The language specified in the full-text policy, \'spa-SPA\', is invalid", e.http_error_message, re.IGNORECASE) async def test_fail_create_full_text_indexing_policy_async(self): From 3901ca37db379d6862d6e50d0fbbeda836fe9e89 Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Thu, 30 Apr 2026 18:29:45 -0400 Subject: [PATCH 10/22] slight improvement on atomicity --- .../tests/test_service_retry_policies.py | 156 ++++------- .../test_service_retry_policies_async.py | 264 ++++++------------ 2 files changed, 138 insertions(+), 282 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/tests/test_service_retry_policies.py b/sdk/cosmos/azure-cosmos/tests/test_service_retry_policies.py index 0a4e19049886..7275d38c662c 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_service_retry_policies.py +++ b/sdk/cosmos/azure-cosmos/tests/test_service_retry_policies.py @@ -2,6 +2,7 @@ # Copyright (c) Microsoft Corporation. All rights reserved. import unittest import uuid +from unittest.mock import patch import pytest from azure.core.exceptions import ServiceRequestError, ServiceResponseError @@ -44,8 +45,6 @@ def test_service_request_retry_policy(self): container = db.get_container_client(self.TEST_CONTAINER_ID) created_item = container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) - # Save the original function - self.original_execute_function = _retry_utility.ExecuteFunction # Change the location cache to have 3 preferred read regions and 3 available read endpoints by location original_location_cache = mock_client.client_connection._global_endpoint_manager.location_cache @@ -57,23 +56,15 @@ def test_service_request_retry_policy(self): self.REGIONAL_ENDPOINT] expected_counter = len(original_location_cache.read_regional_routing_contexts) - try: - # Mock the function to return the ServiceRequestException we retry - mf = self.MockExecuteServiceRequestExceptionIgnoreQuery(self.original_execute_function) - _retry_utility.ExecuteFunction = mf - container.read_item(created_item['id'], created_item['pk']) - pytest.fail("Exception was not raised.") - except ServiceRequestError: - assert mf.counter == expected_counter - finally: - _retry_utility.ExecuteFunction = self.original_execute_function - # Now we test with a query operation, iterating through items sends request without request object - # retry policy should eventually raise an exception as it should stop retrying with a max retry attempt - # equal to the available read region locations + # Test read with IgnoreQuery mock (allows query/pkranges requests through) + mf = self.MockExecuteServiceRequestExceptionIgnoreQuery(_retry_utility.ExecuteFunction) + with patch.object(_retry_utility, 'ExecuteFunction', mf): + with pytest.raises(ServiceRequestError): + container.read_item(created_item['id'], created_item['pk']) + assert mf.counter == expected_counter - # Change the location cache to have 3 preferred read regions and 3 available read endpoints by location - original_location_cache = mock_client.client_connection._global_endpoint_manager.location_cache + # Now we test with a query operation original_location_cache.account_read_locations = [self.REGION1, self.REGION2, self.REGION3] original_location_cache.available_read_regional_endpoints_by_locations = { self.REGION1: self.REGIONAL_ENDPOINT, @@ -83,31 +74,21 @@ def test_service_request_retry_policy(self): self.REGIONAL_ENDPOINT] expected_counter = len(original_location_cache.read_regional_routing_contexts) - try: - # Mock the function to return the ServiceRequestException we retry - mf = self.MockExecuteServiceRequestException() - _retry_utility.ExecuteFunction = mf - items = list(container.query_items(query="SELECT * FROM c", partition_key=created_item['pk'])) - pytest.fail("Exception was not raised.") - except ServiceRequestError: + mf = self.MockExecuteServiceRequestException() + with patch.object(_retry_utility, 'ExecuteFunction', mf): + with pytest.raises(ServiceRequestError): + list(container.query_items(query="SELECT * FROM c", partition_key=created_item['pk'])) assert mf.counter == expected_counter - finally: - _retry_utility.ExecuteFunction = self.original_execute_function # Now we change the location cache to have only 1 preferred read region original_location_cache.account_read_locations = [self.REGION1] original_location_cache.read_regional_routing_contexts = [self.REGIONAL_ENDPOINT] expected_counter = len(original_location_cache.read_regional_routing_contexts) - try: - # Reset the function to reset the counter - mf = self.MockExecuteServiceRequestException() - _retry_utility.ExecuteFunction = mf - container.read_item(created_item['id'], created_item['pk']) - pytest.fail("Exception was not raised.") - except ServiceRequestError: + mf = self.MockExecuteServiceRequestException() + with patch.object(_retry_utility, 'ExecuteFunction', mf): + with pytest.raises(ServiceRequestError): + container.read_item(created_item['id'], created_item['pk']) assert mf.counter == expected_counter - finally: - _retry_utility.ExecuteFunction = self.original_execute_function # Now we try it out with a write request original_location_cache.account_write_locations = [self.REGION1, self.REGION2] @@ -115,17 +96,11 @@ def test_service_request_retry_policy(self): original_location_cache.available_write_regional_endpoints_by_locations = {self.REGION1: self.REGIONAL_ENDPOINT, self.REGION2: self.REGIONAL_ENDPOINT} expected_counter = len(original_location_cache.write_regional_routing_contexts) - try: - # Reset the function to reset the counter - mf = self.MockExecuteServiceRequestException() - _retry_utility.ExecuteFunction = mf - container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) - pytest.fail("Exception was not raised.") - except ServiceRequestError: - # Should retry twice in each region + mf = self.MockExecuteServiceRequestException() + with patch.object(_retry_utility, 'ExecuteFunction', mf): + with pytest.raises(ServiceRequestError): + container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) assert mf.counter == expected_counter - finally: - _retry_utility.ExecuteFunction = self.original_execute_function def test_service_response_retry_policy(self): mock_client = CosmosClient(self.host, self.masterKey) @@ -133,8 +108,6 @@ def test_service_response_retry_policy(self): container = db.get_container_client(self.TEST_CONTAINER_ID) created_item = container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) - # Save the original function - self.original_execute_function = _retry_utility.ExecuteFunction # Change the location cache to have 3 preferred read regions and 3 available read endpoints by location original_location_cache = mock_client.client_connection._global_endpoint_manager.location_cache @@ -144,60 +117,41 @@ def test_service_response_retry_policy(self): self.REGION3: self.REGIONAL_ENDPOINT} original_location_cache.read_regional_routing_contexts = [self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT] - try: - # Mock the function to return the ServiceResponseException we retry - mf = self.MockExecuteServiceResponseExceptionIgnoreQuery(Exception, self.original_execute_function) - _retry_utility.ExecuteFunction = mf - container.read_item(created_item['id'], created_item['pk']) - pytest.fail("Exception was not raised.") - except ServiceResponseError: + + mf = self.MockExecuteServiceResponseExceptionIgnoreQuery(Exception, _retry_utility.ExecuteFunction) + with patch.object(_retry_utility, 'ExecuteFunction', mf): + with pytest.raises(ServiceResponseError): + container.read_item(created_item['id'], created_item['pk']) assert mf.counter == 3 - finally: - _retry_utility.ExecuteFunction = self.original_execute_function # Now we change the location cache to have only 1 preferred read region original_location_cache.account_read_locations = [self.REGION1] original_location_cache.read_regional_routing_contexts = [self.REGIONAL_ENDPOINT] - try: - # Reset the function to reset the counter - mf = self.MockExecuteServiceResponseException(Exception) - _retry_utility.ExecuteFunction = mf - container.read_item(created_item['id'], created_item['pk']) - pytest.fail("Exception was not raised.") - except ServiceResponseError: + mf = self.MockExecuteServiceResponseException(Exception) + with patch.object(_retry_utility, 'ExecuteFunction', mf): + with pytest.raises(ServiceResponseError): + container.read_item(created_item['id'], created_item['pk']) assert mf.counter == 1 - finally: - _retry_utility.ExecuteFunction = self.original_execute_function # Now we try it out with a write request original_location_cache.account_write_locations = [self.REGION1, self.REGION2] original_location_cache.write_regional_routing_contexts = [self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT] original_location_cache.available_write_regional_endpoints_by_locations = {self.REGION1: self.REGIONAL_ENDPOINT, self.REGION2: self.REGIONAL_ENDPOINT} - try: - # Reset the function to reset the counter - mf = self.MockExecuteServiceResponseException(Exception) - _retry_utility.ExecuteFunction = mf - # Even though we have 2 preferred write endpoints, - # we will only run the exception once due to no retries on write requests - container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) - pytest.fail("Exception was not raised.") - except ServiceResponseError: + mf = self.MockExecuteServiceResponseException(Exception) + with patch.object(_retry_utility, 'ExecuteFunction', mf): + with pytest.raises(ServiceResponseError): + # Even though we have 2 preferred write endpoints, + # we will only run the exception once due to no retries on write requests + container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) assert mf.counter == 1 - finally: - _retry_utility.ExecuteFunction = self.original_execute_function # Now we try it out with a write request with retry write enabled - which should retry once - try: - # Reset the function to reset the counter - mf = self.MockExecuteServiceResponseException(Exception) - _retry_utility.ExecuteFunction = mf - container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}, retry_write=2) - pytest.fail("Exception was not raised.") - except ServiceResponseError: + mf = self.MockExecuteServiceResponseException(Exception) + with patch.object(_retry_utility, 'ExecuteFunction', mf): + with pytest.raises(ServiceResponseError): + container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}, retry_write=2) assert mf.counter == 2 - finally: - _retry_utility.ExecuteFunction = self.original_execute_function def test_service_request_connection_retry_policy(self): # Mock the client retry policy to see the same-region retries that happen there @@ -251,38 +205,32 @@ def test_global_endpoint_manager_retry(self): # - GetDatabaseAccountStub allows us to receive any number of endpoints for that call independent of account used exception = ServiceRequestError("mock exception") exception.exc_type = Exception - self.original_get_database_account_stub = _global_endpoint_manager._GlobalEndpointManager._GetDatabaseAccountStub - _global_endpoint_manager._GlobalEndpointManager._GetDatabaseAccountStub = self.MockGetDatabaseAccountStub connection_retry_policy = test_config.MockConnectionRetryPolicy(resource_type="docs", error=exception) - mock_client = CosmosClient(self.host, self.masterKey, connection_retry_policy=connection_retry_policy, - preferred_locations=[self.REGION1, self.REGION2]) - _global_endpoint_manager._GlobalEndpointManager._GetDatabaseAccountStub = self.original_get_database_account_stub + with patch.object(_global_endpoint_manager._GlobalEndpointManager, '_GetDatabaseAccountStub', + self.MockGetDatabaseAccountStub): + mock_client = CosmosClient(self.host, self.masterKey, connection_retry_policy=connection_retry_policy, + preferred_locations=[self.REGION1, self.REGION2]) + db = mock_client.get_database_client(self.TEST_DATABASE_ID) container = db.get_container_client(self.TEST_CONTAINER_ID) - try: - _global_endpoint_manager._GlobalEndpointManager._GetDatabaseAccountStub = self.MockGetDatabaseAccountStub - container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) - pytest.fail("Exception was not raised.") - except ServiceRequestError: + with patch.object(_global_endpoint_manager._GlobalEndpointManager, '_GetDatabaseAccountStub', + self.MockGetDatabaseAccountStub): + with pytest.raises(ServiceRequestError): + container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) assert connection_retry_policy.counter == 3 # 4 total in region retries assert len(connection_retry_policy.request_endpoints) == 4 - finally: - _global_endpoint_manager._GlobalEndpointManager._GetDatabaseAccountStub = self.original_get_database_account_stub # Now we try with a read request - reset the policy to reset the counter connection_retry_policy.request_endpoints = [] - try: - _global_endpoint_manager._GlobalEndpointManager._GetDatabaseAccountStub = self.MockGetDatabaseAccountStub - container.read_item("some_id", "some_pk") - pytest.fail("Exception was not raised.") - except ServiceRequestError: + with patch.object(_global_endpoint_manager._GlobalEndpointManager, '_GetDatabaseAccountStub', + self.MockGetDatabaseAccountStub): + with pytest.raises(ServiceRequestError): + container.read_item("some_id", "some_pk") assert connection_retry_policy.counter == 3 # 4 total requests in each main region (preferred read region 1 -> preferred read region 2) assert len(connection_retry_policy.request_endpoints) == 8 - finally: - _global_endpoint_manager._GlobalEndpointManager._GetDatabaseAccountStub = self.original_get_database_account_stub class MockExecuteServiceRequestException(object): def __init__(self): diff --git a/sdk/cosmos/azure-cosmos/tests/test_service_retry_policies_async.py b/sdk/cosmos/azure-cosmos/tests/test_service_retry_policies_async.py index 7e2235b66b33..e5154e73b506 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_service_retry_policies_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_service_retry_policies_async.py @@ -3,6 +3,7 @@ import unittest import uuid +from unittest.mock import patch import pytest from aiohttp.client_exceptions import (ClientConnectionError, ClientConnectionResetError, @@ -54,8 +55,6 @@ async def test_service_request_retry_policy_async(self): container = db.get_container_client(self.TEST_CONTAINER_ID) created_item = await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) - # Save the original function - self.original_execute_function = _retry_utility_async.ExecuteFunctionAsync # Change the location cache to have 3 preferred read regions and 3 available read endpoints by location original_location_cache = mock_client.client_connection._global_endpoint_manager.location_cache @@ -67,23 +66,14 @@ async def test_service_request_retry_policy_async(self): original_location_cache.read_regional_routing_contexts = [self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT] expected_counter = len(original_location_cache.read_regional_routing_contexts) - try: - # Mock the function to return the ServiceRequestException we retry - mf = self.MockExecuteServiceRequestExceptionIgnoreQuery(self.original_execute_function) - _retry_utility_async.ExecuteFunctionAsync = mf - await container.read_item(created_item['id'], created_item['pk']) - pytest.fail("Exception was not raised.") - except ServiceRequestError: - assert mf.counter == expected_counter - finally: - _retry_utility_async.ExecuteFunctionAsync = self.original_execute_function - # Now we test with a query operation, iterating through items sends request without request object - # retry policy should eventually raise an exception as it should stop retrying with a max retry attempt - # equal to the available read region locations + mf = self.MockExecuteServiceRequestExceptionIgnoreQuery(_retry_utility_async.ExecuteFunctionAsync) + with patch.object(_retry_utility_async, 'ExecuteFunctionAsync', mf): + with pytest.raises(ServiceRequestError): + await container.read_item(created_item['id'], created_item['pk']) + assert mf.counter == expected_counter - # Change the location cache to have 3 preferred read regions and 3 available read endpoints by location - original_location_cache = mock_client.client_connection._global_endpoint_manager.location_cache + # Now we test with a query operation original_location_cache.account_read_locations = [self.REGION1, self.REGION2, self.REGION3] original_location_cache.available_read_regional_endpoints_by_locations = { self.REGION1: self.REGIONAL_ENDPOINT, @@ -93,32 +83,22 @@ async def test_service_request_retry_policy_async(self): self.REGIONAL_ENDPOINT] expected_counter = len(original_location_cache.read_regional_routing_contexts) - try: - # Mock the function to return the ServiceRequestException we retry - mf = self.MockExecuteServiceRequestException() - _retry_utility_async.ExecuteFunctionAsync = mf - items = [item async for item in container.query_items(query="SELECT * FROM c", - partition_key=created_item['pk'])] - pytest.fail("Exception was not raised.") - except ServiceRequestError: + mf = self.MockExecuteServiceRequestException() + with patch.object(_retry_utility_async, 'ExecuteFunctionAsync', mf): + with pytest.raises(ServiceRequestError): + items = [item async for item in container.query_items(query="SELECT * FROM c", + partition_key=created_item['pk'])] assert mf.counter == expected_counter - finally: - _retry_utility_async.ExecuteFunctionAsync = self.original_execute_function # Now we change the location cache to have only 1 preferred read region original_location_cache.account_read_locations = [self.REGION1] original_location_cache.read_regional_routing_contexts = [self.REGIONAL_ENDPOINT] expected_counter = len(original_location_cache.read_regional_routing_contexts) - try: - # Reset the function to reset the counter - mf = self.MockExecuteServiceRequestException() - _retry_utility_async.ExecuteFunctionAsync = mf - await container.read_item(created_item['id'], created_item['pk']) - pytest.fail("Exception was not raised.") - except ServiceRequestError: + mf = self.MockExecuteServiceRequestException() + with patch.object(_retry_utility_async, 'ExecuteFunctionAsync', mf): + with pytest.raises(ServiceRequestError): + await container.read_item(created_item['id'], created_item['pk']) assert mf.counter == expected_counter - finally: - _retry_utility_async.ExecuteFunctionAsync = self.original_execute_function # Now we try it out with a write request original_location_cache.account_write_locations = [self.REGION1, self.REGION2] @@ -127,16 +107,11 @@ async def test_service_request_retry_policy_async(self): self.REGION1: self.REGIONAL_ENDPOINT, self.REGION2: self.REGIONAL_ENDPOINT} expected_counter = len(original_location_cache.write_regional_routing_contexts) - try: - # Reset the function to reset the counter - mf = self.MockExecuteServiceRequestException() - _retry_utility_async.ExecuteFunctionAsync = mf - await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) - pytest.fail("Exception was not raised.") - except ServiceRequestError: + mf = self.MockExecuteServiceRequestException() + with patch.object(_retry_utility_async, 'ExecuteFunctionAsync', mf): + with pytest.raises(ServiceRequestError): + await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) assert mf.counter == expected_counter - finally: - _retry_utility_async.ExecuteFunctionAsync = self.original_execute_function async def test_service_response_retry_policy_async(self): # For ServiceResponseErrors, we only do cross region retries on read requests or on ClientConnectionErrors @@ -146,8 +121,6 @@ async def test_service_response_retry_policy_async(self): container = db.get_container_client(self.TEST_CONTAINER_ID) created_item = await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) - # Save the original function - self.original_execute_function = _retry_utility_async.ExecuteFunctionAsync # Change the location cache to have 3 preferred read regions and 3 available read endpoints by location original_location_cache = mock_client.client_connection._global_endpoint_manager.location_cache @@ -158,31 +131,22 @@ async def test_service_response_retry_policy_async(self): self.REGION3: self.REGIONAL_ENDPOINT} original_location_cache.read_regional_routing_contexts = [self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT] - try: - # Mock the function to return the ClientConnectionError we retry - mf = self.MockExecuteServiceResponseExceptionIgnoreQuery(AttributeError, - None, self.original_execute_function) - _retry_utility_async.ExecuteFunctionAsync = mf - await container.read_item(created_item['id'], created_item['pk']) - pytest.fail("Exception was not raised.") - except ServiceResponseError: + + mf = self.MockExecuteServiceResponseExceptionIgnoreQuery(AttributeError, + None, _retry_utility_async.ExecuteFunctionAsync) + with patch.object(_retry_utility_async, 'ExecuteFunctionAsync', mf): + with pytest.raises(ServiceResponseError): + await container.read_item(created_item['id'], created_item['pk']) assert mf.counter == 3 - finally: - _retry_utility_async.ExecuteFunctionAsync = self.original_execute_function # Now we change the location cache to have only 1 preferred read region original_location_cache.account_read_locations = [self.REGION1] original_location_cache.read_regional_routing_contexts = [self.REGIONAL_ENDPOINT] - try: - # Reset the function to reset the counter - mf = self.MockExecuteServiceResponseException(AttributeError, None) - _retry_utility_async.ExecuteFunctionAsync = mf - await container.read_item(created_item['id'], created_item['pk']) - pytest.fail("Exception was not raised.") - except ServiceResponseError: + mf = self.MockExecuteServiceResponseException(AttributeError, None) + with patch.object(_retry_utility_async, 'ExecuteFunctionAsync', mf): + with pytest.raises(ServiceResponseError): + await container.read_item(created_item['id'], created_item['pk']) assert mf.counter == 1 - finally: - _retry_utility_async.ExecuteFunctionAsync = self.original_execute_function # Now we try it out with a write request original_location_cache.account_write_locations = [self.REGION1, self.REGION2] @@ -190,18 +154,13 @@ async def test_service_response_retry_policy_async(self): original_location_cache.available_write_regional_endpoints_by_locations = { self.REGION1: self.REGIONAL_ENDPOINT, self.REGION2: self.REGIONAL_ENDPOINT} - try: - # Reset the function to reset the counter - mf = self.MockExecuteServiceResponseException(AttributeError, None) - _retry_utility_async.ExecuteFunctionAsync = mf - # Even though we have 2 preferred write endpoints, - # we will only run the exception once due to no retries on write requests - await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) - pytest.fail("Exception was not raised.") - except ServiceResponseError: + mf = self.MockExecuteServiceResponseException(AttributeError, None) + with patch.object(_retry_utility_async, 'ExecuteFunctionAsync', mf): + with pytest.raises(ServiceResponseError): + # Even though we have 2 preferred write endpoints, + # we will only run the exception once due to no retries on write requests + await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) assert mf.counter == 1 - finally: - _retry_utility_async.ExecuteFunctionAsync = self.original_execute_function # If we do a write request with a ClientConnectionError, # we will do cross-region retries like with read requests @@ -210,28 +169,18 @@ async def test_service_response_retry_policy_async(self): original_location_cache.available_write_regional_endpoints_by_locations = { self.REGION1: self.REGIONAL_ENDPOINT, self.REGION2: self.REGIONAL_ENDPOINT} - try: - # Reset the function to reset the counter - mf = self.MockExecuteServiceResponseException(ClientConnectionError, ClientConnectionError()) - _retry_utility_async.ExecuteFunctionAsync = mf - await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) - pytest.fail("Exception was not raised.") - except ServiceResponseError: + mf = self.MockExecuteServiceResponseException(ClientConnectionError, ClientConnectionError()) + with patch.object(_retry_utility_async, 'ExecuteFunctionAsync', mf): + with pytest.raises(ServiceResponseError): + await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) assert mf.counter == 2 - finally: - _retry_utility_async.ExecuteFunctionAsync = self.original_execute_function # Now we try it out with a write request with retry write enabled - which should retry once - try: - # Reset the function to reset the counter - mf = self.MockExecuteServiceResponseException(AttributeError, None) - _retry_utility_async.ExecuteFunctionAsync = mf - await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}, retry_write=2) - pytest.fail("Exception was not raised.") - except ServiceResponseError: + mf = self.MockExecuteServiceResponseException(AttributeError, None) + with patch.object(_retry_utility_async, 'ExecuteFunctionAsync', mf): + with pytest.raises(ServiceResponseError): + await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}, retry_write=2) assert mf.counter == 2 - finally: - _retry_utility_async.ExecuteFunctionAsync = self.original_execute_function async def test_service_request_connection_retry_policy_async(self): # Mock the client retry policy to see the same-region retries that happen there @@ -285,8 +234,6 @@ async def test_service_response_connection_retry_policy_async(self): async def test_service_response_errors_async(self): # Test for errors that are subclasses of ClientConnectionError for write requests - # Save the original ExecuteAsyncFunction function - self.original_execute_function = _retry_utility_async.ExecuteFunctionAsync async with CosmosClient(self.host, self.masterKey) as mock_client: db = mock_client.get_database_client(self.TEST_DATABASE_ID) container = db.get_container_client(self.TEST_CONTAINER_ID) @@ -305,105 +252,77 @@ async def test_service_response_errors_async(self): original_location_cache.available_write_regional_endpoints_by_locations = { self.REGION1: self.REGIONAL_ENDPOINT, self.REGION2: self.REGIONAL_ENDPOINT} - try: - # Start with a normal ServiceResponseException with no special casing - mf = self.MockExecuteServiceResponseException(AttributeError, AttributeError()) - _retry_utility_async.ExecuteFunctionAsync = mf - # Even though we have 2 preferred write endpoints, - # we will only run the exception once due to no retries on write requests - await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) - pytest.fail("Exception was not raised.") - except ServiceResponseError: + + # Start with a normal ServiceResponseException with no special casing + mf = self.MockExecuteServiceResponseException(AttributeError, AttributeError()) + with patch.object(_retry_utility_async, 'ExecuteFunctionAsync', mf): + with pytest.raises(ServiceResponseError): + # Even though we have 2 preferred write endpoints, + # we will only run the exception once due to no retries on write requests + await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) assert mf.counter == 1 - finally: - _retry_utility_async.ExecuteFunctionAsync = self.original_execute_function # Now we test the base ClientConnectionError to see in-region retry - try: - # Reset the function to reset the counter - mf = self.MockExecuteServiceResponseException(ClientConnectionError, ClientConnectionError()) - _retry_utility_async.ExecuteFunctionAsync = mf - await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) - pytest.fail("Exception was not raised.") - except ServiceResponseError: + mf = self.MockExecuteServiceResponseException(ClientConnectionError, ClientConnectionError()) + with patch.object(_retry_utility_async, 'ExecuteFunctionAsync', mf): + with pytest.raises(ServiceResponseError): + await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) assert mf.counter == 2 assert len(original_location_cache.location_unavailability_info_by_endpoint) == 1 host_unavailable = original_location_cache.location_unavailability_info_by_endpoint.get(self.host) assert host_unavailable is not None assert len(host_unavailable.get('operationType')) == 1 assert 'Write' in host_unavailable.get('operationType') - finally: - _retry_utility_async.ExecuteFunctionAsync = self.original_execute_function await container.read() # We send another request to the same error - since we marked one of the two in-region endpoints unavailable we don't retry - try: - # Reset the function to reset the counter - mf = self.MockExecuteServiceResponseException(ClientConnectionError, ClientConnectionError()) - _retry_utility_async.ExecuteFunctionAsync = mf - await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) - pytest.fail("Exception was not raised.") - except ServiceResponseError: + mf = self.MockExecuteServiceResponseException(ClientConnectionError, ClientConnectionError()) + with patch.object(_retry_utility_async, 'ExecuteFunctionAsync', mf): + with pytest.raises(ServiceResponseError): + await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) assert mf.counter == 1 host_unavailable = original_location_cache.location_unavailability_info_by_endpoint.get(self.host) assert host_unavailable is not None assert len(host_unavailable.get('operationType')) == 1 assert 'Write' in host_unavailable.get('operationType') - finally: - _retry_utility_async.ExecuteFunctionAsync = self.original_execute_function # Reset the location cache's unavailable endpoints in order to try the same with other exceptions original_location_cache.location_unavailability_info_by_endpoint = {} original_location_cache.write_regional_routing_contexts = [self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT] # Now we test ClientConnectionResetError, the subclass of ClientConnectionError - try: - # Reset the function to reset the counter - mf = self.MockExecuteServiceResponseException(ClientConnectionResetError, ClientConnectionResetError()) - _retry_utility_async.ExecuteFunctionAsync = mf - await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) - pytest.fail("Exception was not raised.") - except ServiceResponseError: + mf = self.MockExecuteServiceResponseException(ClientConnectionResetError, ClientConnectionResetError()) + with patch.object(_retry_utility_async, 'ExecuteFunctionAsync', mf): + with pytest.raises(ServiceResponseError): + await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) assert mf.counter == 2 assert len(original_location_cache.location_unavailability_info_by_endpoint) == 1 host_unavailable = original_location_cache.location_unavailability_info_by_endpoint.get(self.host) assert host_unavailable is not None assert len(host_unavailable.get('operationType')) == 1 assert 'Write' in host_unavailable.get('operationType') - finally: - _retry_utility_async.ExecuteFunctionAsync = self.original_execute_function # Reset the location cache's unavailable endpoints in order to try the same with other exceptions original_location_cache.location_unavailability_info_by_endpoint = {} original_location_cache.write_regional_routing_contexts = [self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT] # Now we test ServerConnectionError, the subclass of ClientConnectionError - try: - # Reset the function to reset the counter - mf = self.MockExecuteServiceResponseException(ServerConnectionError, ServerConnectionError()) - _retry_utility_async.ExecuteFunctionAsync = mf - await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) - pytest.fail("Exception was not raised.") - except ServiceResponseError: + mf = self.MockExecuteServiceResponseException(ServerConnectionError, ServerConnectionError()) + with patch.object(_retry_utility_async, 'ExecuteFunctionAsync', mf): + with pytest.raises(ServiceResponseError): + await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) assert mf.counter == 2 - finally: - _retry_utility_async.ExecuteFunctionAsync = self.original_execute_function # Reset the location cache's unavailable endpoints in order to try the same with other exceptions original_location_cache.location_unavailability_info_by_endpoint = {} original_location_cache.write_regional_routing_contexts = [self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT] # Now we test ClientOSError, the subclass of ClientConnectionError - try: - # Reset the function to reset the counter - mf = self.MockExecuteServiceResponseException(ClientOSError, ClientOSError()) - _retry_utility_async.ExecuteFunctionAsync = mf - await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) - pytest.fail("Exception was not raised.") - except ServiceResponseError: + mf = self.MockExecuteServiceResponseException(ClientOSError, ClientOSError()) + with patch.object(_retry_utility_async, 'ExecuteFunctionAsync', mf): + with pytest.raises(ServiceResponseError): + await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) assert mf.counter == 2 - finally: - _retry_utility_async.ExecuteFunctionAsync = self.original_execute_function async def test_global_endpoint_manager_retry_async(self): # For this test we mock both the ConnectionRetryPolicy and the GetDatabaseAccountStub @@ -411,40 +330,29 @@ async def test_global_endpoint_manager_retry_async(self): # - GetDatabaseAccountStub allows us to receive any number of endpoints for that call independent of account used exception = ServiceRequestError("mock exception") exception.exc_type = Exception - self.original_get_database_account_stub = _global_endpoint_manager_async._GlobalEndpointManager._GetDatabaseAccountStub - _global_endpoint_manager_async._GlobalEndpointManager._GetDatabaseAccountStub = self.MockGetDatabaseAccountStub connection_policy = self.connectionPolicy connection_retry_policy = test_config.MockConnectionRetryPolicyAsync(resource_type="docs", error=exception) connection_policy.ConnectionRetryConfiguration = connection_retry_policy - async with CosmosClient(self.host, self.masterKey, connection_policy=connection_policy, - preferred_locations=[self.REGION1, self.REGION2]) as mock_client: - db = mock_client.get_database_client(self.TEST_DATABASE_ID) - container = db.get_container_client(self.TEST_CONTAINER_ID) - - try: - await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) - pytest.fail("Exception was not raised.") - except ServiceRequestError: + with patch.object(_global_endpoint_manager_async._GlobalEndpointManager, '_GetDatabaseAccountStub', + self.MockGetDatabaseAccountStub): + async with CosmosClient(self.host, self.masterKey, connection_policy=connection_policy, + preferred_locations=[self.REGION1, self.REGION2]) as mock_client: + db = mock_client.get_database_client(self.TEST_DATABASE_ID) + container = db.get_container_client(self.TEST_CONTAINER_ID) + + with pytest.raises((ServiceRequestError, CosmosHttpResponseError)): + await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) assert connection_retry_policy.counter == 3 # 4 total in region retries assert len(connection_retry_policy.request_endpoints) == 4 - except CosmosHttpResponseError as e: - print(e) - finally: - _global_endpoint_manager_async._GlobalEndpointManager._GetDatabaseAccountStub = self.original_get_database_account_stub - - # Now we try with a read request - reset the policy to reset the counter - _global_endpoint_manager_async._GlobalEndpointManager._GetDatabaseAccountStub = self.MockGetDatabaseAccountStub - connection_retry_policy.request_endpoints = [] - try: - await container.read_item("some_id", "some_pk") - pytest.fail("Exception was not raised.") - except ServiceRequestError: + + # Now we try with a read request - reset the policy to reset the counter + connection_retry_policy.request_endpoints = [] + with pytest.raises(ServiceRequestError): + await container.read_item("some_id", "some_pk") assert connection_retry_policy.counter == 3 # 4 total requests in each main region (preferred read region 1 -> preferred read region 2) assert len(connection_retry_policy.request_endpoints) == 8 - finally: - _global_endpoint_manager_async._GlobalEndpointManager._GetDatabaseAccountStub = self.original_get_database_account_stub class MockExecuteServiceRequestException(object): def __init__(self): From 61672375ee3fbdcf1203bb305f4d0e68ad5b65fc Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Fri, 1 May 2026 17:31:35 -0400 Subject: [PATCH 11/22] Fix retry policy test race condition with background health check Cancel the GEM's background _endpoints_health_check task before mutating cache state in async tests. Set all underlying source-of-truth attributes (account_*_regional_routing_contexts_by_location, effective_preferred_locations) consistently via helper methods, so any recalculation by update_location_cache() produces expected results. Remove references to non-existent available_read/write_regional_endpoints_by_locations attribute. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../tests/test_service_retry_policies.py | 57 +++++----- .../test_service_retry_policies_async.py | 101 +++++++++--------- 2 files changed, 73 insertions(+), 85 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/tests/test_service_retry_policies.py b/sdk/cosmos/azure-cosmos/tests/test_service_retry_policies.py index 7275d38c662c..0a3555a908cf 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_service_retry_policies.py +++ b/sdk/cosmos/azure-cosmos/tests/test_service_retry_policies.py @@ -39,6 +39,21 @@ def setUpClass(cls): cls.created_database = cls.client.get_database_client(cls.TEST_DATABASE_ID) cls.created_container = cls.created_database.get_container_client(cls.TEST_CONTAINER_ID) + def _setup_read_regions(self, location_cache, regions): + """Set all read region attributes consistently so update_location_cache() recalculates correctly.""" + location_cache.account_read_locations = regions + location_cache.account_read_regional_routing_contexts_by_location = { + r: self.REGIONAL_ENDPOINT for r in regions} + location_cache.read_regional_routing_contexts = [self.REGIONAL_ENDPOINT] * len(regions) + location_cache.effective_preferred_locations = regions + + def _setup_write_regions(self, location_cache, regions): + """Set all write region attributes consistently so update_location_cache() recalculates correctly.""" + location_cache.account_write_locations = regions + location_cache.account_write_regional_routing_contexts_by_location = { + r: self.REGIONAL_ENDPOINT for r in regions} + location_cache.write_regional_routing_contexts = [self.REGIONAL_ENDPOINT] * len(regions) + def test_service_request_retry_policy(self): mock_client = CosmosClient(self.host, self.masterKey) db = mock_client.get_database_client(self.TEST_DATABASE_ID) @@ -46,14 +61,9 @@ def test_service_request_retry_policy(self): created_item = container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) - # Change the location cache to have 3 preferred read regions and 3 available read endpoints by location + # Change the location cache to have 3 preferred read regions original_location_cache = mock_client.client_connection._global_endpoint_manager.location_cache - original_location_cache.account_read_locations = [self.REGION1, self.REGION2, self.REGION3] - original_location_cache.available_read_regional_endpoints_by_locations = {self.REGION1: self.REGIONAL_ENDPOINT, - self.REGION2: self.REGIONAL_ENDPOINT, - self.REGION3: self.REGIONAL_ENDPOINT} - original_location_cache.read_regional_routing_contexts = [self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT, - self.REGIONAL_ENDPOINT] + self._setup_read_regions(original_location_cache, [self.REGION1, self.REGION2, self.REGION3]) expected_counter = len(original_location_cache.read_regional_routing_contexts) @@ -65,13 +75,7 @@ def test_service_request_retry_policy(self): assert mf.counter == expected_counter # Now we test with a query operation - original_location_cache.account_read_locations = [self.REGION1, self.REGION2, self.REGION3] - original_location_cache.available_read_regional_endpoints_by_locations = { - self.REGION1: self.REGIONAL_ENDPOINT, - self.REGION2: self.REGIONAL_ENDPOINT, - self.REGION3: self.REGIONAL_ENDPOINT} - original_location_cache.read_regional_routing_contexts = [self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT, - self.REGIONAL_ENDPOINT] + self._setup_read_regions(original_location_cache, [self.REGION1, self.REGION2, self.REGION3]) expected_counter = len(original_location_cache.read_regional_routing_contexts) mf = self.MockExecuteServiceRequestException() @@ -81,8 +85,7 @@ def test_service_request_retry_policy(self): assert mf.counter == expected_counter # Now we change the location cache to have only 1 preferred read region - original_location_cache.account_read_locations = [self.REGION1] - original_location_cache.read_regional_routing_contexts = [self.REGIONAL_ENDPOINT] + self._setup_read_regions(original_location_cache, [self.REGION1]) expected_counter = len(original_location_cache.read_regional_routing_contexts) mf = self.MockExecuteServiceRequestException() with patch.object(_retry_utility, 'ExecuteFunction', mf): @@ -91,10 +94,7 @@ def test_service_request_retry_policy(self): assert mf.counter == expected_counter # Now we try it out with a write request - original_location_cache.account_write_locations = [self.REGION1, self.REGION2] - original_location_cache.write_regional_routing_contexts = [self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT] - original_location_cache.available_write_regional_endpoints_by_locations = {self.REGION1: self.REGIONAL_ENDPOINT, - self.REGION2: self.REGIONAL_ENDPOINT} + self._setup_write_regions(original_location_cache, [self.REGION1, self.REGION2]) expected_counter = len(original_location_cache.write_regional_routing_contexts) mf = self.MockExecuteServiceRequestException() with patch.object(_retry_utility, 'ExecuteFunction', mf): @@ -109,14 +109,9 @@ def test_service_response_retry_policy(self): created_item = container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) - # Change the location cache to have 3 preferred read regions and 3 available read endpoints by location + # Change the location cache to have 3 preferred read regions original_location_cache = mock_client.client_connection._global_endpoint_manager.location_cache - original_location_cache.account_read_locations = [self.REGION1, self.REGION2, self.REGION3] - original_location_cache.available_read_regional_endpoints_by_locations = {self.REGION1: self.REGIONAL_ENDPOINT, - self.REGION2: self.REGIONAL_ENDPOINT, - self.REGION3: self.REGIONAL_ENDPOINT} - original_location_cache.read_regional_routing_contexts = [self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT, - self.REGIONAL_ENDPOINT] + self._setup_read_regions(original_location_cache, [self.REGION1, self.REGION2, self.REGION3]) mf = self.MockExecuteServiceResponseExceptionIgnoreQuery(Exception, _retry_utility.ExecuteFunction) with patch.object(_retry_utility, 'ExecuteFunction', mf): @@ -125,8 +120,7 @@ def test_service_response_retry_policy(self): assert mf.counter == 3 # Now we change the location cache to have only 1 preferred read region - original_location_cache.account_read_locations = [self.REGION1] - original_location_cache.read_regional_routing_contexts = [self.REGIONAL_ENDPOINT] + self._setup_read_regions(original_location_cache, [self.REGION1]) mf = self.MockExecuteServiceResponseException(Exception) with patch.object(_retry_utility, 'ExecuteFunction', mf): with pytest.raises(ServiceResponseError): @@ -134,10 +128,7 @@ def test_service_response_retry_policy(self): assert mf.counter == 1 # Now we try it out with a write request - original_location_cache.account_write_locations = [self.REGION1, self.REGION2] - original_location_cache.write_regional_routing_contexts = [self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT] - original_location_cache.available_write_regional_endpoints_by_locations = {self.REGION1: self.REGIONAL_ENDPOINT, - self.REGION2: self.REGIONAL_ENDPOINT} + self._setup_write_regions(original_location_cache, [self.REGION1, self.REGION2]) mf = self.MockExecuteServiceResponseException(Exception) with patch.object(_retry_utility, 'ExecuteFunction', mf): with pytest.raises(ServiceResponseError): diff --git a/sdk/cosmos/azure-cosmos/tests/test_service_retry_policies_async.py b/sdk/cosmos/azure-cosmos/tests/test_service_retry_policies_async.py index e5154e73b506..92419c71fb34 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_service_retry_policies_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_service_retry_policies_async.py @@ -48,6 +48,33 @@ async def asyncTearDown(self): self.connectionPolicy.ConnectionRetryConfiguration = None await self.client.close() + async def _cancel_background_refresh_task(self, client): + """Cancel the GEM's background health check task to prevent it from + overwriting test cache mutations via update_location_cache().""" + gem = client.client_connection._global_endpoint_manager + if gem.refresh_task and not gem.refresh_task.done(): + gem.refresh_task.cancel() + try: + await gem.refresh_task + except BaseException: + pass + gem.refresh_task = None + + def _setup_read_regions(self, location_cache, regions): + """Set all read region attributes consistently so update_location_cache() recalculates correctly.""" + location_cache.account_read_locations = regions + location_cache.account_read_regional_routing_contexts_by_location = { + r: self.REGIONAL_ENDPOINT for r in regions} + location_cache.read_regional_routing_contexts = [self.REGIONAL_ENDPOINT] * len(regions) + location_cache.effective_preferred_locations = regions + + def _setup_write_regions(self, location_cache, regions): + """Set all write region attributes consistently so update_location_cache() recalculates correctly.""" + location_cache.account_write_locations = regions + location_cache.account_write_regional_routing_contexts_by_location = { + r: self.REGIONAL_ENDPOINT for r in regions} + location_cache.write_regional_routing_contexts = [self.REGIONAL_ENDPOINT] * len(regions) + async def test_service_request_retry_policy_async(self): # ServiceRequestErrors will always retry, and will retry once per preferred region async with CosmosClient(self.host, self.masterKey) as mock_client: @@ -56,15 +83,12 @@ async def test_service_request_retry_policy_async(self): created_item = await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) - # Change the location cache to have 3 preferred read regions and 3 available read endpoints by location + # Cancel background health check to prevent it from overwriting cache mutations + await self._cancel_background_refresh_task(mock_client) + + # Change the location cache to have 3 preferred read regions original_location_cache = mock_client.client_connection._global_endpoint_manager.location_cache - original_location_cache.account_read_locations = [self.REGION1, self.REGION2, self.REGION3] - original_location_cache.available_read_regional_endpoints_by_locations = { - self.REGION1: self.REGIONAL_ENDPOINT, - self.REGION2: self.REGIONAL_ENDPOINT, - self.REGION3: self.REGIONAL_ENDPOINT} - original_location_cache.read_regional_routing_contexts = [self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT, - self.REGIONAL_ENDPOINT] + self._setup_read_regions(original_location_cache, [self.REGION1, self.REGION2, self.REGION3]) expected_counter = len(original_location_cache.read_regional_routing_contexts) mf = self.MockExecuteServiceRequestExceptionIgnoreQuery(_retry_utility_async.ExecuteFunctionAsync) @@ -74,13 +98,7 @@ async def test_service_request_retry_policy_async(self): assert mf.counter == expected_counter # Now we test with a query operation - original_location_cache.account_read_locations = [self.REGION1, self.REGION2, self.REGION3] - original_location_cache.available_read_regional_endpoints_by_locations = { - self.REGION1: self.REGIONAL_ENDPOINT, - self.REGION2: self.REGIONAL_ENDPOINT, - self.REGION3: self.REGIONAL_ENDPOINT} - original_location_cache.read_regional_routing_contexts = [self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT, - self.REGIONAL_ENDPOINT] + self._setup_read_regions(original_location_cache, [self.REGION1, self.REGION2, self.REGION3]) expected_counter = len(original_location_cache.read_regional_routing_contexts) mf = self.MockExecuteServiceRequestException() @@ -91,8 +109,7 @@ async def test_service_request_retry_policy_async(self): assert mf.counter == expected_counter # Now we change the location cache to have only 1 preferred read region - original_location_cache.account_read_locations = [self.REGION1] - original_location_cache.read_regional_routing_contexts = [self.REGIONAL_ENDPOINT] + self._setup_read_regions(original_location_cache, [self.REGION1]) expected_counter = len(original_location_cache.read_regional_routing_contexts) mf = self.MockExecuteServiceRequestException() with patch.object(_retry_utility_async, 'ExecuteFunctionAsync', mf): @@ -101,11 +118,7 @@ async def test_service_request_retry_policy_async(self): assert mf.counter == expected_counter # Now we try it out with a write request - original_location_cache.account_write_locations = [self.REGION1, self.REGION2] - original_location_cache.write_regional_routing_contexts = [self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT] - original_location_cache.available_write_regional_endpoints_by_locations = { - self.REGION1: self.REGIONAL_ENDPOINT, - self.REGION2: self.REGIONAL_ENDPOINT} + self._setup_write_regions(original_location_cache, [self.REGION1, self.REGION2]) expected_counter = len(original_location_cache.write_regional_routing_contexts) mf = self.MockExecuteServiceRequestException() with patch.object(_retry_utility_async, 'ExecuteFunctionAsync', mf): @@ -122,15 +135,12 @@ async def test_service_response_retry_policy_async(self): created_item = await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) - # Change the location cache to have 3 preferred read regions and 3 available read endpoints by location + # Cancel background health check to prevent it from overwriting cache mutations + await self._cancel_background_refresh_task(mock_client) + + # Change the location cache to have 3 preferred read regions original_location_cache = mock_client.client_connection._global_endpoint_manager.location_cache - original_location_cache.account_read_locations = [self.REGION1, self.REGION2, self.REGION3] - original_location_cache.available_read_regional_endpoints_by_locations = { - self.REGION1: self.REGIONAL_ENDPOINT, - self.REGION2: self.REGIONAL_ENDPOINT, - self.REGION3: self.REGIONAL_ENDPOINT} - original_location_cache.read_regional_routing_contexts = [self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT, - self.REGIONAL_ENDPOINT] + self._setup_read_regions(original_location_cache, [self.REGION1, self.REGION2, self.REGION3]) mf = self.MockExecuteServiceResponseExceptionIgnoreQuery(AttributeError, None, _retry_utility_async.ExecuteFunctionAsync) @@ -140,8 +150,7 @@ async def test_service_response_retry_policy_async(self): assert mf.counter == 3 # Now we change the location cache to have only 1 preferred read region - original_location_cache.account_read_locations = [self.REGION1] - original_location_cache.read_regional_routing_contexts = [self.REGIONAL_ENDPOINT] + self._setup_read_regions(original_location_cache, [self.REGION1]) mf = self.MockExecuteServiceResponseException(AttributeError, None) with patch.object(_retry_utility_async, 'ExecuteFunctionAsync', mf): with pytest.raises(ServiceResponseError): @@ -149,11 +158,7 @@ async def test_service_response_retry_policy_async(self): assert mf.counter == 1 # Now we try it out with a write request - original_location_cache.account_write_locations = [self.REGION1, self.REGION2] - original_location_cache.write_regional_routing_contexts = [self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT] - original_location_cache.available_write_regional_endpoints_by_locations = { - self.REGION1: self.REGIONAL_ENDPOINT, - self.REGION2: self.REGIONAL_ENDPOINT} + self._setup_write_regions(original_location_cache, [self.REGION1, self.REGION2]) mf = self.MockExecuteServiceResponseException(AttributeError, None) with patch.object(_retry_utility_async, 'ExecuteFunctionAsync', mf): with pytest.raises(ServiceResponseError): @@ -164,11 +169,7 @@ async def test_service_response_retry_policy_async(self): # If we do a write request with a ClientConnectionError, # we will do cross-region retries like with read requests - original_location_cache.account_write_locations = [self.REGION1, self.REGION2] - original_location_cache.write_regional_routing_contexts = [self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT] - original_location_cache.available_write_regional_endpoints_by_locations = { - self.REGION1: self.REGIONAL_ENDPOINT, - self.REGION2: self.REGIONAL_ENDPOINT} + self._setup_write_regions(original_location_cache, [self.REGION1, self.REGION2]) mf = self.MockExecuteServiceResponseException(ClientConnectionError, ClientConnectionError()) with patch.object(_retry_utility_async, 'ExecuteFunctionAsync', mf): with pytest.raises(ServiceResponseError): @@ -239,19 +240,15 @@ async def test_service_response_errors_async(self): container = db.get_container_client(self.TEST_CONTAINER_ID) await container.read() + # Cancel background health check to prevent it from overwriting cache mutations + await self._cancel_background_refresh_task(mock_client) + original_location_cache = mock_client.client_connection._global_endpoint_manager.location_cache - original_location_cache.account_read_locations = [self.REGION1, self.REGION2, self.REGION3] - original_location_cache.available_read_regional_endpoints_by_locations = { - self.REGION1: self.REGIONAL_ENDPOINT, - self.REGION2: self.REGIONAL_ENDPOINT, - self.REGION3: self.REGIONAL_ENDPOINT} - original_location_cache.read_regional_routing_contexts = [self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT, - self.REGIONAL_ENDPOINT] + self._setup_read_regions(original_location_cache, [self.REGION1, self.REGION2, self.REGION3]) + # For writes, set only the derived state directly since test_service_response_errors + # relies on mark_endpoint_unavailable -> update_location_cache() reducing write regions original_location_cache.account_write_locations = [self.REGION1, self.REGION2] original_location_cache.write_regional_routing_contexts = [self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT] - original_location_cache.available_write_regional_endpoints_by_locations = { - self.REGION1: self.REGIONAL_ENDPOINT, - self.REGION2: self.REGIONAL_ENDPOINT} # Start with a normal ServiceResponseException with no special casing mf = self.MockExecuteServiceResponseException(AttributeError, AttributeError()) From ae88b5bd245337fa4ed732cd25b330a6d2e574aa Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Fri, 1 May 2026 17:47:26 -0400 Subject: [PATCH 12/22] Add cache consistency and health check interaction tests Add 3 unit tests to test_location_cache.py: - test_update_location_cache_recalculates_from_account_data - test_update_location_cache_preserves_unavailability_marks - test_location_cache_derived_state_consistency Add 2 integration tests to test_health_check_async.py: - test_health_check_overwrites_stale_cache_mutations_async - test_retry_policy_unaffected_by_concurrent_cache_refresh_async These tests document the cache architecture invariants that caused the retry test CI failures and serve as regression tests for future changes to the health check or location cache. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../tests/test_health_check_async.py | 133 ++++++++++++++++++ .../azure-cosmos/tests/test_location_cache.py | 115 +++++++++++++++ 2 files changed, 248 insertions(+) diff --git a/sdk/cosmos/azure-cosmos/tests/test_health_check_async.py b/sdk/cosmos/azure-cosmos/tests/test_health_check_async.py index 165dcb151d73..a405686f040a 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_health_check_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_health_check_async.py @@ -4,6 +4,7 @@ import asyncio import time import unittest +import unittest.mock import uuid from typing import List from azure.cosmos.exceptions import CosmosHttpResponseError @@ -427,5 +428,137 @@ async def tracking_mock(self_gem, endpoint, **kwargs): _global_endpoint_manager_async._GlobalEndpointManager._GetDatabaseAccountStub = self.original_getDatabaseAccountStub + async def test_health_check_overwrites_stale_cache_mutations_async(self, setup): + """ + Verify that the background health check task overwrites any stale direct + mutations to routing contexts. This documents WHY tests that manipulate + cache state must cancel the background task first. + """ + from azure.cosmos._location_cache import RegionalRoutingContext + + self.original_getDatabaseAccountStub = _global_endpoint_manager_async._GlobalEndpointManager._GetDatabaseAccountStub + _global_endpoint_manager_async._GlobalEndpointManager._GetDatabaseAccountStub = ( + self.MockGetDatabaseAccount(REGIONS)) + + try: + client = CosmosClient(self.host, self.masterKey, preferred_locations=REGIONS) + await client.__aenter__() + gem = client.client_connection._global_endpoint_manager + location_cache = gem.location_cache + + # Wait for initial health check to complete + await asyncio.sleep(2) + + # Now directly mutate the read_regional_routing_contexts (simulating what old tests did) + fake_contexts = [ + RegionalRoutingContext("https://fake-region1.documents.azure.com"), + RegionalRoutingContext("https://fake-region2.documents.azure.com"), + ] + location_cache.read_regional_routing_contexts = fake_contexts + + # Verify our mutation took effect + assert location_cache.read_regional_routing_contexts == fake_contexts + + # Trigger a refresh that will run update_location_cache + gem.refresh_needed = True + gem.startup = False + + # Do an operation to trigger the refresh cycle + self.original_health_check_method = _cosmos_client_connection_async.CosmosClientConnection.health_check + mock_probe = self.MockHealthCheckProbe() + _cosmos_client_connection_async.CosmosClientConnection.health_check = mock_probe + + # Force a cache update (simulates what the health check does on completion) + location_cache.update_location_cache() + + # The stale mutation should be overwritten by recalculation from account data + current_contexts = location_cache.read_regional_routing_contexts + current_endpoints = [ctx.get_primary() for ctx in current_contexts] + assert "https://fake-region1.documents.azure.com" not in current_endpoints, \ + "Stale direct mutations should be overwritten by update_location_cache()" + assert "https://fake-region2.documents.azure.com" not in current_endpoints, \ + "Stale direct mutations should be overwritten by update_location_cache()" + + _cosmos_client_connection_async.CosmosClientConnection.health_check = self.original_health_check_method + await client.close() + finally: + _global_endpoint_manager_async._GlobalEndpointManager._GetDatabaseAccountStub = self.original_getDatabaseAccountStub + + async def test_retry_policy_unaffected_by_concurrent_cache_refresh_async(self, setup): + """ + Verify that the retry policy's failover_retry_count is unaffected by + a mid-flight update_location_cache() call. The retry counter lives on + the policy object, not in the location cache. + """ + from azure.cosmos._service_request_retry_policy import ServiceRequestRetryPolicy + from azure.cosmos._location_cache import RegionalRoutingContext + from azure.cosmos._request_object import RequestObject + from azure.cosmos.documents import _OperationType + from azure.cosmos.http_constants import ResourceType + + self.original_getDatabaseAccountStub = _global_endpoint_manager_async._GlobalEndpointManager._GetDatabaseAccountStub + _global_endpoint_manager_async._GlobalEndpointManager._GetDatabaseAccountStub = ( + self.MockGetDatabaseAccount(REGIONS)) + + try: + client = CosmosClient(self.host, self.masterKey, preferred_locations=REGIONS) + await client.__aenter__() + gem = client.client_connection._global_endpoint_manager + location_cache = gem.location_cache + + # Wait for initial setup + await asyncio.sleep(2) + + # Cancel background task to have deterministic control + if hasattr(gem, 'refresh_task') and gem.refresh_task and not gem.refresh_task.done(): + gem.refresh_task.cancel() + try: + await gem.refresh_task + except BaseException: + pass + + # Setup: ensure we have multiple read regions + read_contexts = location_cache.read_regional_routing_contexts + assert len(read_contexts) >= 2, "Need at least 2 read regions for this test" + + # Create a retry policy + request = RequestObject(ResourceType.Document, _OperationType.Read, None) + request.location_endpoint_to_route = read_contexts[0].get_primary() + + mock_gem = unittest.mock.Mock() + mock_gem.location_cache = location_cache + mock_gem.resolve_service_endpoint_for_partition.return_value = read_contexts[1].get_primary() + mock_gem.mark_endpoint_unavailable_for_read = location_cache.mark_endpoint_unavailable_for_read + + mock_connection_policy = unittest.mock.Mock() + mock_connection_policy.EnableEndpointDiscovery = True + mock_pk_range_wrapper = unittest.mock.Mock() + + retry_policy = ServiceRequestRetryPolicy( + mock_connection_policy, mock_gem, mock_pk_range_wrapper, request) + + # First retry should succeed + should_retry = retry_policy.ShouldRetry() + assert should_retry is True + assert retry_policy.failover_retry_count == 1 + + # Now simulate a concurrent cache refresh mid-flight + location_cache.update_location_cache() + + # The retry count should be unaffected by the cache refresh + assert retry_policy.failover_retry_count == 1, \ + "Retry count should not be affected by concurrent cache refresh" + + # Second retry attempt — count should increment normally + mock_gem.resolve_service_endpoint_for_partition.return_value = read_contexts[0].get_primary() + should_retry_again = retry_policy.ShouldRetry() + assert retry_policy.failover_retry_count == 2, \ + "Retry count should increment normally regardless of cache refresh" + + await client.close() + finally: + _global_endpoint_manager_async._GlobalEndpointManager._GetDatabaseAccountStub = self.original_getDatabaseAccountStub + + if __name__ == '__main__': unittest.main() diff --git a/sdk/cosmos/azure-cosmos/tests/test_location_cache.py b/sdk/cosmos/azure-cosmos/tests/test_location_cache.py index 8f44511cfdb9..b94280e7febc 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_location_cache.py +++ b/sdk/cosmos/azure-cosmos/tests/test_location_cache.py @@ -555,5 +555,120 @@ def test_unavailable_endpoints_ordering_in_routing_list(self): assert write_contexts[1].get_primary() == location1_endpoint assert write_contexts[2].get_primary() == location2_endpoint + def test_update_location_cache_recalculates_from_account_data(self): + """ + Verify that update_location_cache() recalculates routing contexts from the + underlying account data (account_*_regional_routing_contexts_by_location), + not from stale direct mutations to read/write_regional_routing_contexts. + + This documents the invariant that caused the retry test CI failures: + if you mutate read_regional_routing_contexts directly without updating + the source-of-truth attributes, a subsequent update_location_cache() call + will overwrite your mutations. + """ + preferred_locations = [location1_name, location2_name, location4_name] + lc = refresh_location_cache(preferred_locations, use_multiple_write_locations=False) + db_acc = create_database_account(enable_multiple_writable_locations=False) + lc.perform_on_database_account_read(db_acc) + + # Capture the correctly-derived state + original_read_contexts = list(lc.read_regional_routing_contexts) + assert len(original_read_contexts) == 3 # location1, location2, location4 + + # Simulate what the old tests did: directly mutate the derived list + from azure.cosmos._location_cache import RegionalRoutingContext + lc.read_regional_routing_contexts = [ + RegionalRoutingContext("https://fake1.documents.azure.com"), + RegionalRoutingContext("https://fake2.documents.azure.com"), + RegionalRoutingContext("https://fake3.documents.azure.com"), + ] + + # Now call update_location_cache() with no new data (simulates health check completion) + lc.update_location_cache() + + # The direct mutation should be overwritten — routing contexts recalculated from account data + recalculated = lc.read_regional_routing_contexts + assert len(recalculated) == len(original_read_contexts) + for orig, recalc in zip(original_read_contexts, recalculated): + assert orig.get_primary() == recalc.get_primary(), \ + "update_location_cache() should recalculate from account data, overwriting direct mutations" + + def test_update_location_cache_preserves_unavailability_marks(self): + """ + Verify that mark_endpoint_unavailable() state survives subsequent + update_location_cache() calls. Unavailable endpoints should be moved + to the end of the routing list, not removed. + """ + preferred_locations = [location1_name, location2_name, location3_name] + lc = refresh_location_cache(preferred_locations, use_multiple_write_locations=True) + db_acc = create_database_account(enable_multiple_writable_locations=True) + lc.perform_on_database_account_read(db_acc) + + # Mark location1 as unavailable for reads + lc.mark_endpoint_unavailable_for_read(location1_endpoint, refresh_cache=False) + + # Verify the unavailability is recorded + assert lc.is_endpoint_unavailable(location1_endpoint, "Read") + + # Call update_location_cache() — simulates background health check completing + lc.update_location_cache() + + # Unavailability info should persist (it lives in location_unavailability_info_by_endpoint) + assert lc.is_endpoint_unavailable(location1_endpoint, "Read"), \ + "Unavailability marks should survive update_location_cache() calls" + + # The unavailable endpoint should be moved to the end of the read routing list + read_contexts = lc.get_read_regional_routing_contexts() + read_endpoints = [ctx.get_primary() for ctx in read_contexts] + assert location1_endpoint in read_endpoints, "Unavailable endpoint should still be in routing list" + assert read_endpoints[-1] == location1_endpoint, \ + "Unavailable endpoint should be at the end of routing list" + assert read_endpoints[0] == location2_endpoint, \ + "First healthy preferred endpoint should be first" + + def test_location_cache_derived_state_consistency(self): + """ + Verify that update_location_cache() is idempotent and that the derived + state (read/write_regional_routing_contexts) matches what + get_preferred_regional_routing_contexts() would return independently. + """ + preferred_locations = [location1_name, location2_name, location4_name] + lc = refresh_location_cache(preferred_locations, use_multiple_write_locations=True) + db_acc = create_database_account(enable_multiple_writable_locations=True) + lc.perform_on_database_account_read(db_acc) + + # Snapshot after first update + read_after_first = [ctx.get_primary() for ctx in lc.read_regional_routing_contexts] + write_after_first = [ctx.get_primary() for ctx in lc.write_regional_routing_contexts] + + # Call update_location_cache() again with no new data — should be idempotent + lc.update_location_cache() + + read_after_second = [ctx.get_primary() for ctx in lc.read_regional_routing_contexts] + write_after_second = [ctx.get_primary() for ctx in lc.write_regional_routing_contexts] + + assert read_after_first == read_after_second, \ + "update_location_cache() should be idempotent for read routing contexts" + assert write_after_first == write_after_second, \ + "update_location_cache() should be idempotent for write routing contexts" + + # Verify derived state matches what get_preferred_regional_routing_contexts returns directly + from azure.cosmos._location_cache import EndpointOperationType + expected_read = lc.get_preferred_regional_routing_contexts( + lc.account_read_regional_routing_contexts_by_location, + lc.account_read_locations, + EndpointOperationType.ReadType, + lc.write_regional_routing_contexts[0] + ) + expected_write = lc.get_preferred_regional_routing_contexts( + lc.account_write_regional_routing_contexts_by_location, + lc.account_write_locations, + EndpointOperationType.WriteType, + lc.default_regional_routing_context + ) + + assert read_after_second == [ctx.get_primary() for ctx in expected_read] + assert write_after_second == [ctx.get_primary() for ctx in expected_write] + if __name__ == "__main__": unittest.main() From fbcf5827edf2226ae4d505c5ed5b2b77da062e9b Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Fri, 1 May 2026 18:58:37 -0400 Subject: [PATCH 13/22] Update test_location_cache.py --- sdk/cosmos/azure-cosmos/tests/test_location_cache.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/tests/test_location_cache.py b/sdk/cosmos/azure-cosmos/tests/test_location_cache.py index b94280e7febc..e68736d78b5a 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_location_cache.py +++ b/sdk/cosmos/azure-cosmos/tests/test_location_cache.py @@ -589,8 +589,8 @@ def test_update_location_cache_recalculates_from_account_data(self): # The direct mutation should be overwritten — routing contexts recalculated from account data recalculated = lc.read_regional_routing_contexts assert len(recalculated) == len(original_read_contexts) - for orig, recalc in zip(original_read_contexts, recalculated): - assert orig.get_primary() == recalc.get_primary(), \ + for orig, context in zip(original_read_contexts, recalculated): + assert orig.get_primary() == context.get_primary(), \ "update_location_cache() should recalculate from account data, overwriting direct mutations" def test_update_location_cache_preserves_unavailability_marks(self): From bc8a023173852319849edfe145230c9562180e68 Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Sat, 2 May 2026 14:16:19 -0400 Subject: [PATCH 14/22] stop asserting error text - service keeps changing it --- sdk/cosmos/azure-cosmos/tests/test_full_text_policy.py | 6 ------ .../azure-cosmos/tests/test_full_text_policy_async.py | 8 +------- 2 files changed, 1 insertion(+), 13 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/tests/test_full_text_policy.py b/sdk/cosmos/azure-cosmos/tests/test_full_text_policy.py index c70906844aca..e2260fe28091 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_full_text_policy.py +++ b/sdk/cosmos/azure-cosmos/tests/test_full_text_policy.py @@ -230,9 +230,6 @@ def test_fail_create_full_text_policy(self): pytest.fail("Container creation should have failed for wrong supported language.") except exceptions.CosmosHttpResponseError as e: assert e.status_code == 400 - assert re.search( - r"The language specified in the full-text policy, \'spa-SPA\', is invalid.", - e.http_error_message, re.IGNORECASE) # Pass a full text policy with an unsupported path language full_text_policy_wrong_default = { @@ -253,9 +250,6 @@ def test_fail_create_full_text_policy(self): pytest.fail("Container creation should have failed for wrong supported language.") except exceptions.CosmosHttpResponseError as e: assert e.status_code == 400 - assert re.search( - r"The language specified in the full-text policy, \'spa-SPA\', is invalid", - e.http_error_message, re.IGNORECASE) def test_fail_create_full_text_indexing_policy(self): full_text_policy = { diff --git a/sdk/cosmos/azure-cosmos/tests/test_full_text_policy_async.py b/sdk/cosmos/azure-cosmos/tests/test_full_text_policy_async.py index c4b67f07f2ef..2dd51e83faf1 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_full_text_policy_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_full_text_policy_async.py @@ -199,7 +199,7 @@ async def test_replace_full_text_container_async(self): assert properties["indexingPolicy"]['fullTextIndexes'] == indexing_policy['fullTextIndexes'] assert created_container_properties['fullTextPolicy'] != properties['fullTextPolicy'] assert created_container_properties["indexingPolicy"] != properties["indexingPolicy"] - self.test_db.delete_container(created_container.id) + await self.test_db.delete_container(created_container.id) async def test_fail_create_full_text_policy_async(self): # Pass a full text policy with a wrongly formatted path @@ -244,9 +244,6 @@ async def test_fail_create_full_text_policy_async(self): pytest.fail("Container creation should have failed for wrong supported language.") except exceptions.CosmosHttpResponseError as e: assert e.status_code == 400 - assert re.search( - r"The language specified in the full-text policy, \'spa-SPA\', is invalid.", - e.http_error_message, re.IGNORECASE) # Pass a full text policy with an unsupported path language full_text_policy_wrong_default = { @@ -267,9 +264,6 @@ async def test_fail_create_full_text_policy_async(self): pytest.fail("Container creation should have failed for wrong supported language.") except exceptions.CosmosHttpResponseError as e: assert e.status_code == 400 - assert re.search( - r"The language specified in the full-text policy, \'spa-SPA\', is invalid", - e.http_error_message, re.IGNORECASE) async def test_fail_create_full_text_indexing_policy_async(self): full_text_policy = { From 7eeebf7a0850204f1340dfbe345e7087b26656de Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Sat, 2 May 2026 18:29:49 -0400 Subject: [PATCH 15/22] Update test_computed_properties_async.py --- .../azure-cosmos/tests/test_computed_properties_async.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/sdk/cosmos/azure-cosmos/tests/test_computed_properties_async.py b/sdk/cosmos/azure-cosmos/tests/test_computed_properties_async.py index 3b4a35a97318..e93b873add82 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_computed_properties_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_computed_properties_async.py @@ -1,6 +1,7 @@ # The MIT License (MIT) # Copyright (c) Microsoft Corporation. All rights reserved. +import asyncio import unittest import uuid import pytest @@ -177,6 +178,9 @@ async def test_replace_with_new_computed_properties_async(self): computed_properties=new_computed_properties ) + # Allow time for the emulator to reindex existing items with new computed properties + await asyncio.sleep(2) + # Check if computed properties were set container = await replaced_collection.read() assert new_computed_properties == container["computedProperties"] From 6351fe14538a41926ed49744a37c8b2bc25cef2a Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Sat, 2 May 2026 22:14:38 -0400 Subject: [PATCH 16/22] Update test_computed_properties_async.py --- .../azure-cosmos/tests/test_computed_properties_async.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/tests/test_computed_properties_async.py b/sdk/cosmos/azure-cosmos/tests/test_computed_properties_async.py index e93b873add82..e5b91be1fb0f 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_computed_properties_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_computed_properties_async.py @@ -178,8 +178,9 @@ async def test_replace_with_new_computed_properties_async(self): computed_properties=new_computed_properties ) - # Allow time for the emulator to reindex existing items with new computed properties - await asyncio.sleep(2) + # Re-upsert items so computed properties are indexed with the new definitions + for item in self.items: + await replaced_collection.upsert_item(body=item) # Check if computed properties were set container = await replaced_collection.read() From f66518dcbcc568147f7aa1ade912040a05c5aab5 Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Sun, 3 May 2026 13:25:03 -0400 Subject: [PATCH 17/22] Revert computed properties async test to original state The test_replace_with_new_computed_properties_async failure (0 == 3) is a pre-existing emulator timing issue, not a regression from this PR. Reverting our attempted fixes (asyncio.sleep and upsert) since they did not resolve the underlying issue. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../azure-cosmos/tests/test_computed_properties_async.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/tests/test_computed_properties_async.py b/sdk/cosmos/azure-cosmos/tests/test_computed_properties_async.py index e5b91be1fb0f..3b4a35a97318 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_computed_properties_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_computed_properties_async.py @@ -1,7 +1,6 @@ # The MIT License (MIT) # Copyright (c) Microsoft Corporation. All rights reserved. -import asyncio import unittest import uuid import pytest @@ -178,10 +177,6 @@ async def test_replace_with_new_computed_properties_async(self): computed_properties=new_computed_properties ) - # Re-upsert items so computed properties are indexed with the new definitions - for item in self.items: - await replaced_collection.upsert_item(body=item) - # Check if computed properties were set container = await replaced_collection.read() assert new_computed_properties == container["computedProperties"] From 085cf2628cc7f95dfee99b684d88a6f7f46b05e2 Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Sun, 3 May 2026 14:42:05 -0400 Subject: [PATCH 18/22] Add diagnostic + polling to computed properties async test Adds two diagnostic tools to test_replace_with_new_computed_properties_async to investigate why it returns 0 items instead of 3 after replace_container: 1. Diagnostic query: SELECT c.id, c.db_group, c.cp_upper FROM c to inspect whether items exist and whether cp_upper is being computed at all. 2. Polling retry: retries the cp_upper WHERE query up to 10 times with 1-second delays to determine if this is an emulator reindexing lag issue. Both print [DIAG] output visible in CI logs. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../tests/test_computed_properties_async.py | 28 +++++++++++++++++-- 1 file changed, 25 insertions(+), 3 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/tests/test_computed_properties_async.py b/sdk/cosmos/azure-cosmos/tests/test_computed_properties_async.py index 3b4a35a97318..500a3cf3135d 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_computed_properties_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_computed_properties_async.py @@ -1,6 +1,7 @@ # The MIT License (MIT) # Copyright (c) Microsoft Corporation. All rights reserved. +import asyncio import unittest import uuid import pytest @@ -181,10 +182,31 @@ async def test_replace_with_new_computed_properties_async(self): container = await replaced_collection.read() assert new_computed_properties == container["computedProperties"] + # Diagnostic: query all items and their computed property values to understand indexing state + all_items = [q async for q in + replaced_collection.query_items( + query='SELECT c.id, c.db_group, c.cp_upper FROM c', + partition_key="test")] + print(f"\n[DIAG] Total items after replace: {len(all_items)}") + for diag_item in all_items[:3]: + print(f"[DIAG] Item id={diag_item.get('id', '?')}, " + f"db_group={diag_item.get('db_group', '?')}, " + f"cp_upper={diag_item.get('cp_upper', 'MISSING')}") + + # Polling: retry the cp_upper query with delays to detect emulator reindexing lag + queried_items = [] + for attempt in range(10): + queried_items = [q async for q in + replaced_collection.query_items( + query='Select * from c Where c.cp_upper = "GROUP2"', + partition_key="test")] + if len(queried_items) == 3: + print(f"[DIAG] cp_upper query succeeded on attempt {attempt + 1}") + break + print(f"[DIAG] cp_upper query attempt {attempt + 1}: got {len(queried_items)}, expected 3") + await asyncio.sleep(1) + # Test 1: Test first computed property - queried_items = [q async for q in - replaced_collection.query_items(query='Select * from c Where c.cp_upper = "GROUP2"', - partition_key="test")] self.assertEqual(len(queried_items), 3) # Test 1 Negative: Test if using non-existent computed property name returns nothing From d0dbaa1cd1e7ebc34a2d4671d46153ad40712541 Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Sun, 3 May 2026 16:32:00 -0400 Subject: [PATCH 19/22] Update cspell.json --- sdk/cosmos/azure-cosmos/cspell.json | 1 + 1 file changed, 1 insertion(+) diff --git a/sdk/cosmos/azure-cosmos/cspell.json b/sdk/cosmos/azure-cosmos/cspell.json index fd2c6923e829..913606d7e7ff 100644 --- a/sdk/cosmos/azure-cosmos/cspell.json +++ b/sdk/cosmos/azure-cosmos/cspell.json @@ -6,6 +6,7 @@ "perfresults", "pkrange", "ppcb", + "reindexing", "reranker", "toctou" ] From b8292b3749e78eb833980062061ad9b0767c5eed Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Mon, 4 May 2026 16:51:04 -0400 Subject: [PATCH 20/22] Address review comments: rename split test, fix finally block - Renamed test_full_refresh_fallback_after_split_preserves_containerRID to test_incremental_retry_after_split_preserves_containerRID (sync + async). The test actually validates the incremental retry path, not the full refresh fallback. Updated docstrings and comments to match. - Moved health_check monkey-patch revert into the finally block in test_health_check_overwrites_stale_cache_mutations_async to prevent leaking into other tests if an assertion fails mid-test. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../tests/test_container_rid_header_unit.py | 16 +++++++++------- .../test_container_rid_header_unit_async.py | 16 +++++++++------- .../tests/test_health_check_async.py | 3 ++- 3 files changed, 20 insertions(+), 15 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/tests/test_container_rid_header_unit.py b/sdk/cosmos/azure-cosmos/tests/test_container_rid_header_unit.py index 8b155a44fbcf..4e7cc871d782 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_container_rid_header_unit.py +++ b/sdk/cosmos/azure-cosmos/tests/test_container_rid_header_unit.py @@ -171,10 +171,11 @@ def test_incremental_refresh_preserves_containerRID(self): assert client.call_count == 2 assert client.captured_feed_options.get("containerRID") == CONTAINER_RID - def test_full_refresh_fallback_after_split_preserves_containerRID(self): + def test_incremental_retry_after_split_preserves_containerRID(self): """When an incremental refresh encounters unresolvable parent ranges, - it falls back to a full refresh. containerRID must still flow through - on the fallback call.""" + the routing map provider retries the incremental fetch once + (_INCOMPLETE_ROUTING_MAP_MAX_RETRIES=1). On the retry, if the ranges + resolve successfully, containerRID must still flow through.""" # Post-split ranges reference parents not in the original map split_ranges = [ {"id": "3", "minInclusive": "", "maxExclusive": "1F", "parents": ["99"]}, @@ -190,8 +191,9 @@ def _ReadPartitionKeyRanges(self, _collection_link, feed_options=None, **kwargs) self.captured_feed_options = dict(feed_options) if feed_options else {} self.call_count += 1 # First call: initial load returns original ranges - # Second call: incremental returns split ranges with unknown parents - # Third call: full fallback returns original ranges + # Second call: incremental returns split ranges with unknown parents, + # triggering _NeedFullRefresh → retry (count 0→1) + # Third call: incremental retry succeeds with resolvable ranges if self.call_count == 2: return iter(split_ranges) return iter(PARTITION_KEY_RANGES) @@ -208,8 +210,8 @@ def _ReadPartitionKeyRanges(self, _collection_link, feed_options=None, **kwargs) force_refresh=True, previous_routing_map=previous_map, ) - # Incremental fetch returns ranges with unresolvable parents, triggering - # a full refresh fallback: 1 initial + 1 incremental + 1 full = 3 calls. + # Incremental fetch hits unresolvable parents → retries once → succeeds: + # 1 initial + 1 failed incremental + 1 successful retry = 3 calls. assert client.call_count == 3 assert client.captured_feed_options.get("containerRID") == CONTAINER_RID diff --git a/sdk/cosmos/azure-cosmos/tests/test_container_rid_header_unit_async.py b/sdk/cosmos/azure-cosmos/tests/test_container_rid_header_unit_async.py index 5de1cfbb9379..73e45b7ebdbe 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_container_rid_header_unit_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_container_rid_header_unit_async.py @@ -161,10 +161,11 @@ async def test_incremental_refresh_preserves_containerRID_async(self): assert client.call_count == 2 assert client.captured_feed_options.get("containerRID") == CONTAINER_RID - async def test_full_refresh_fallback_after_split_preserves_containerRID_async(self): + async def test_incremental_retry_after_split_preserves_containerRID_async(self): """When an incremental refresh encounters unresolvable parent ranges, - it falls back to a full refresh. containerRID must still flow through - on the fallback call.""" + the routing map provider retries the incremental fetch once + (_INCOMPLETE_ROUTING_MAP_MAX_RETRIES=1). On the retry, if the ranges + resolve successfully, containerRID must still flow through.""" # Post-split ranges reference parents not in the original map split_ranges = [ {"id": "3", "minInclusive": "", "maxExclusive": "1F", "parents": ["99"]}, @@ -188,8 +189,9 @@ async def _ReadPartitionKeyRanges(self, _collection_link, feed_options=None, **k self.call_count += 1 SplitMockClient._capture_internal_headers(kwargs, '"test-etag-1"') # First call: initial load returns original ranges - # Second call: incremental returns split ranges with unknown parents - # Third call: full fallback returns original ranges + # Second call: incremental returns split ranges with unknown parents, + # triggering _NeedFullRefresh → retry (count 0→1) + # Third call: incremental retry succeeds with resolvable ranges items = split_ranges if self.call_count == 2 else PARTITION_KEY_RANGES for item in items: yield item @@ -206,8 +208,8 @@ async def _ReadPartitionKeyRanges(self, _collection_link, feed_options=None, **k force_refresh=True, previous_routing_map=previous_map, ) - # Incremental fetch returns ranges with unresolvable parents, triggering - # a full refresh fallback: 1 initial + 1 incremental + 1 full = 3 calls. + # Incremental fetch hits unresolvable parents → retries once → succeeds: + # 1 initial + 1 failed incremental + 1 successful retry = 3 calls. assert client.call_count == 3 assert client.captured_feed_options.get("containerRID") == CONTAINER_RID diff --git a/sdk/cosmos/azure-cosmos/tests/test_health_check_async.py b/sdk/cosmos/azure-cosmos/tests/test_health_check_async.py index a405686f040a..2eb7a15c2580 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_health_check_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_health_check_async.py @@ -479,10 +479,11 @@ async def test_health_check_overwrites_stale_cache_mutations_async(self, setup): assert "https://fake-region2.documents.azure.com" not in current_endpoints, \ "Stale direct mutations should be overwritten by update_location_cache()" - _cosmos_client_connection_async.CosmosClientConnection.health_check = self.original_health_check_method await client.close() finally: _global_endpoint_manager_async._GlobalEndpointManager._GetDatabaseAccountStub = self.original_getDatabaseAccountStub + if hasattr(self, 'original_health_check_method'): + _cosmos_client_connection_async.CosmosClientConnection.health_check = self.original_health_check_method async def test_retry_policy_unaffected_by_concurrent_cache_refresh_async(self, setup): """ From b4d7b6d5996366f69992eb98c8f48bc71e858abf Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Mon, 4 May 2026 17:09:41 -0400 Subject: [PATCH 21/22] Add full refresh fallback test for containerRID preservation Tests the path where incremental retries are exhausted and the routing map provider falls back to a full refresh (previous_map=None), verifying containerRID still flows through on the full refresh call (4 total calls). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../tests/test_container_rid_header_unit.py | 43 ++++++++++++++++ .../test_container_rid_header_unit_async.py | 51 +++++++++++++++++++ 2 files changed, 94 insertions(+) diff --git a/sdk/cosmos/azure-cosmos/tests/test_container_rid_header_unit.py b/sdk/cosmos/azure-cosmos/tests/test_container_rid_header_unit.py index 4e7cc871d782..5572029608b0 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_container_rid_header_unit.py +++ b/sdk/cosmos/azure-cosmos/tests/test_container_rid_header_unit.py @@ -215,6 +215,49 @@ def _ReadPartitionKeyRanges(self, _collection_link, feed_options=None, **kwargs) assert client.call_count == 3 assert client.captured_feed_options.get("containerRID") == CONTAINER_RID + def test_full_refresh_fallback_after_exhausted_retries_preserves_containerRID(self): + """When incremental retries are exhausted (_INCOMPLETE_ROUTING_MAP_MAX_RETRIES=1), + the routing map provider falls back to a full refresh (current_previous_map=None). + containerRID must still flow through on the full refresh call.""" + # Post-split ranges reference parents not in the original map + split_ranges = [ + {"id": "3", "minInclusive": "", "maxExclusive": "1F", "parents": ["99"]}, + {"id": "4", "minInclusive": "1F", "maxExclusive": "3F", "parents": ["99"]}, + ] + + class FullRefreshMockClient: + def __init__(self): + self.captured_feed_options = None + self.call_count = 0 + + def _ReadPartitionKeyRanges(self, _collection_link, feed_options=None, **kwargs): + self.captured_feed_options = dict(feed_options) if feed_options else {} + self.call_count += 1 + # Call 1: initial load → good ranges + # Call 2: incremental → split ranges (unresolvable) → retry (count 0→1) + # Call 3: incremental retry → split ranges again (still unresolvable) + # → retries exhausted → fall back to full refresh (previous_map=None) + # Call 4: full refresh → good ranges → success + if self.call_count in (2, 3): + return iter(split_ranges) + return iter(PARTITION_KEY_RANGES) + + client = FullRefreshMockClient() + cache = PartitionKeyRangeCache(client) + feed_options = {"containerRID": CONTAINER_RID} + previous_map = cache.get_routing_map(COLLECTION_LINK, feed_options) + assert client.call_count == 1 + assert client.captured_feed_options.get("containerRID") == CONTAINER_RID + cache.get_routing_map( + COLLECTION_LINK, + feed_options, + force_refresh=True, + previous_routing_map=previous_map, + ) + # 1 initial + 1 failed incremental + 1 failed retry + 1 full refresh = 4 calls. + assert client.call_count == 4 + assert client.captured_feed_options.get("containerRID") == CONTAINER_RID + # ----- Cache staleness detection ----- def _make_cache_with_map(self, collection_id, etag): diff --git a/sdk/cosmos/azure-cosmos/tests/test_container_rid_header_unit_async.py b/sdk/cosmos/azure-cosmos/tests/test_container_rid_header_unit_async.py index 73e45b7ebdbe..5f6196bd4d5f 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_container_rid_header_unit_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_container_rid_header_unit_async.py @@ -213,6 +213,57 @@ async def _ReadPartitionKeyRanges(self, _collection_link, feed_options=None, **k assert client.call_count == 3 assert client.captured_feed_options.get("containerRID") == CONTAINER_RID + async def test_full_refresh_fallback_after_exhausted_retries_preserves_containerRID_async(self): + """When incremental retries are exhausted (_INCOMPLETE_ROUTING_MAP_MAX_RETRIES=1), + the routing map provider falls back to a full refresh (current_previous_map=None). + containerRID must still flow through on the full refresh call.""" + # Post-split ranges reference parents not in the original map + split_ranges = [ + {"id": "3", "minInclusive": "", "maxExclusive": "1F", "parents": ["99"]}, + {"id": "4", "minInclusive": "1F", "maxExclusive": "3F", "parents": ["99"]}, + ] + + class FullRefreshMockClient: + def __init__(self): + self.captured_feed_options = None + self.call_count = 0 + + @staticmethod + def _capture_internal_headers(kwargs, etag): + captured_headers = kwargs.get('_internal_response_headers_capture') + if captured_headers is not None: + captured_headers.clear() + captured_headers.update({'ETag': etag}) + + async def _ReadPartitionKeyRanges(self, _collection_link, feed_options=None, **kwargs): + self.captured_feed_options = dict(feed_options) if feed_options else {} + self.call_count += 1 + FullRefreshMockClient._capture_internal_headers(kwargs, '"test-etag-1"') + # Call 1: initial load → good ranges + # Call 2: incremental → split ranges (unresolvable) → retry (count 0→1) + # Call 3: incremental retry → split ranges again (still unresolvable) + # → retries exhausted → fall back to full refresh (previous_map=None) + # Call 4: full refresh → good ranges → success + items = split_ranges if self.call_count in (2, 3) else PARTITION_KEY_RANGES + for item in items: + yield item + + client = FullRefreshMockClient() + cache = PartitionKeyRangeCache(client) + feed_options: Dict[str, Any] = {"containerRID": CONTAINER_RID} + previous_map = await cache.get_routing_map(COLLECTION_LINK, feed_options) + assert client.call_count == 1 + assert client.captured_feed_options.get("containerRID") == CONTAINER_RID + await cache.get_routing_map( + COLLECTION_LINK, + feed_options, + force_refresh=True, + previous_routing_map=previous_map, + ) + # 1 initial + 1 failed incremental + 1 failed retry + 1 full refresh = 4 calls. + assert client.call_count == 4 + assert client.captured_feed_options.get("containerRID") == CONTAINER_RID + async def test_full_load_removes_stale_if_none_match_header_async(self): """When async _fetch_routing_map falls back to a full load (no previous From 468c9cdd81a60ce3c02b4d1cc962d568feb77308 Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Mon, 4 May 2026 18:12:58 -0400 Subject: [PATCH 22/22] rename error to reflect actual behavior --- .../_routing/_routing_map_provider_common.py | 18 ++++++++++-------- .../_routing/aio/routing_map_provider.py | 4 ++-- .../cosmos/_routing/routing_map_provider.py | 4 ++-- .../tests/test_container_rid_header_unit.py | 2 +- .../test_container_rid_header_unit_async.py | 2 +- 5 files changed, 16 insertions(+), 14 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/_routing_map_provider_common.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/_routing_map_provider_common.py index 60a8b224b37a..ce579fdb258a 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/_routing_map_provider_common.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/_routing_map_provider_common.py @@ -148,9 +148,12 @@ def _resolve_endpoint(client: Any) -> str: return f"__unknown_{id(client)}__" -class _NeedFullRefresh(Exception): +class _IncrementalMergeFailed(Exception): """Sentinel raised by :func:`process_fetched_ranges` when the - incremental update cannot be completed and a full refresh is needed.""" + incremental update cannot resolve all partition key ranges. + + The caller decides how to recover: retry the incremental fetch + (if attempts remain) or fall back to a full routing-map refresh.""" def process_fetched_ranges( @@ -178,10 +181,9 @@ def process_fetched_ranges( initial load yields no ranges. :rtype: ~azure.cosmos._routing.collection_routing_map.CollectionRoutingMap or None - :raises _NeedFullRefresh: When the incremental path determines a - full refresh is required. The caller should catch this and - re-invoke ``_fetch_routing_map`` with - ``previous_routing_map=None``. + :raises _IncrementalMergeFailed: When the incremental path cannot + resolve all ranges. The caller catches this and either retries + the incremental fetch or falls back to a full refresh. """ if not previous_routing_map: # Initial load -- build the complete map. @@ -251,7 +253,7 @@ def process_fetched_ranges( collection_link, first_unresolved.get(PartitionKeyRange.Id), ) - raise _NeedFullRefresh() + raise _IncrementalMergeFailed() unresolved = next_unresolved @@ -262,7 +264,7 @@ def process_fetched_ranges( "collection '%s'. Falling back to full refresh.", collection_link, ) - raise _NeedFullRefresh() + raise _IncrementalMergeFailed() return result diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/aio/routing_map_provider.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/aio/routing_map_provider.py index 009d999a31d8..4cfb429ab7e3 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/aio/routing_map_provider.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/aio/routing_map_provider.py @@ -37,7 +37,7 @@ is_cache_unchanged_since_previous, determine_refresh_action, get_smart_overlapping_ranges, - _NeedFullRefresh, + _IncrementalMergeFailed, ) @@ -375,7 +375,7 @@ async def _fetch_routing_map( return process_fetched_ranges( ranges, current_previous_map, collection_id, collection_link, new_etag ) - except _NeedFullRefresh: + except _IncrementalMergeFailed: if current_previous_map is not None and incomplete_attempt_count < _INCOMPLETE_ROUTING_MAP_MAX_RETRIES: incomplete_attempt_count += 1 logger.warning( diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_map_provider.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_map_provider.py index be65f0a128e6..92abe54cba94 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_map_provider.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_map_provider.py @@ -36,7 +36,7 @@ is_cache_unchanged_since_previous, determine_refresh_action, get_smart_overlapping_ranges, - _NeedFullRefresh, + _IncrementalMergeFailed, ) if TYPE_CHECKING: @@ -345,7 +345,7 @@ def _fetch_routing_map( return process_fetched_ranges( ranges, current_previous_map, collection_id, collection_link, new_etag ) - except _NeedFullRefresh: + except _IncrementalMergeFailed: if current_previous_map is not None and incomplete_attempt_count < _INCOMPLETE_ROUTING_MAP_MAX_RETRIES: incomplete_attempt_count += 1 logger.warning( diff --git a/sdk/cosmos/azure-cosmos/tests/test_container_rid_header_unit.py b/sdk/cosmos/azure-cosmos/tests/test_container_rid_header_unit.py index 5572029608b0..8ea0bda6dbe3 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_container_rid_header_unit.py +++ b/sdk/cosmos/azure-cosmos/tests/test_container_rid_header_unit.py @@ -192,7 +192,7 @@ def _ReadPartitionKeyRanges(self, _collection_link, feed_options=None, **kwargs) self.call_count += 1 # First call: initial load returns original ranges # Second call: incremental returns split ranges with unknown parents, - # triggering _NeedFullRefresh → retry (count 0→1) + # raising _IncrementalMergeFailed (caught → retry incremental, count 0→1) # Third call: incremental retry succeeds with resolvable ranges if self.call_count == 2: return iter(split_ranges) diff --git a/sdk/cosmos/azure-cosmos/tests/test_container_rid_header_unit_async.py b/sdk/cosmos/azure-cosmos/tests/test_container_rid_header_unit_async.py index 5f6196bd4d5f..85b90785fbfc 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_container_rid_header_unit_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_container_rid_header_unit_async.py @@ -190,7 +190,7 @@ async def _ReadPartitionKeyRanges(self, _collection_link, feed_options=None, **k SplitMockClient._capture_internal_headers(kwargs, '"test-etag-1"') # First call: initial load returns original ranges # Second call: incremental returns split ranges with unknown parents, - # triggering _NeedFullRefresh → retry (count 0→1) + # raising _IncrementalMergeFailed (caught → retry incremental, count 0→1) # Third call: incremental retry succeeds with resolvable ranges items = split_ranges if self.call_count == 2 else PARTITION_KEY_RANGES for item in items: