# The MIT License (MIT) # Copyright (c) Microsoft Corporation. All rights reserved. import pytest import unittest import uuid from azure.core.exceptions import ServiceResponseError from azure.core.rest import HttpRequest from azure.core.pipeline.transport._requests_basic import RequestsTransport from azure.cosmos import CosmosClient, ContainerProxy, DatabaseProxy, _retry_utility, exceptions from azure.cosmos.partition_key import PartitionKey from _fault_injection_transport import FaultInjectionTransport from typing import Callable import test_config # Import test_config for configuration details HPK_CONTAINER = "TestContainerHPK" @pytest.fixture(scope="class") def setup(): config = test_config.TestConfig() if config.masterKey == '[YOUR_KEY_HERE]' or config.host == '[YOUR_ENDPOINT_HERE]': raise Exception( "You must specify your Azure Cosmos account values for " "'masterKey' and 'host' at the top of this class to run the " "tests.") test_client = CosmosClient(config.host, config.masterKey) database = test_client.create_database_if_not_exists(id=config.TEST_DATABASE_ID) created_container = database.create_container_if_not_exists( id=config.TEST_SINGLE_PARTITION_CONTAINER_ID, partition_key=PartitionKey(path="/pk") ) container_hpk = database.create_container_if_not_exists( id=HPK_CONTAINER, partition_key=PartitionKey(path=["/state", "/city"], kind="MultiHash") ) yield {"database": database, "container": created_container, "container_hpk": container_hpk} @pytest.mark.cosmosEmulator @pytest.mark.usefixtures("setup") class TestRetryableWrites: host = test_config.TestConfig.host master_key = test_config.TestConfig.masterKey TEST_DATABASE_ID = test_config.TestConfig.TEST_DATABASE_ID TEST_CONTAINER_SINGLE_PARTITION_ID = test_config.TestConfig.TEST_SINGLE_PARTITION_CONTAINER_ID TEST_CONTAINER_MULTI_PARTITION_ID = test_config.TestConfig.TEST_MULTI_PARTITION_CONTAINER_ID client: CosmosClient = None @staticmethod def setup_method_with_custom_transport( custom_transport: RequestsTransport, default_endpoint: str = host, key: str = master_key, database_id: str = TEST_DATABASE_ID, container_id: str = TEST_CONTAINER_SINGLE_PARTITION_ID, **kwargs): client = CosmosClient(default_endpoint, key, consistency_level="Session", transport=custom_transport, **kwargs) db: DatabaseProxy = client.create_database_if_not_exists(database_id) container: ContainerProxy = db.create_container_if_not_exists(container_id, PartitionKey(path="/pk")) return {"client": client, "db": db, "col": container} def test_retryable_writes_request_level(self, setup): # Create a container for testing container = setup["container"] # Mock retry_utility.execute to track retries original_execute = _retry_utility.ExecuteFunction me = self.MockExecuteFunction(original_execute) _retry_utility.ExecuteFunction = me # Test without retry_write for upsert_item test_item = {"id": "1", "pk": "test", "data": "retryable write test"} try: container.upsert_item(test_item) pytest.fail("Expected an exception without retry_write") except exceptions.CosmosHttpResponseError: assert me.counter == 1, "Expected no retries without retry_write" # Test with retry_write for upsert_item me.counter = 0 # Reset counter try: container.upsert_item(test_item, retry_write=2) except exceptions.CosmosHttpResponseError as e: assert me.counter > 1, "Expected multiple retries due to simulated errors" finally: # Restore the original retry_utility.execute function _retry_utility.ExecuteFunction = original_execute # Verify the item was written read_item = container.read_item(item=test_item['id'], partition_key=test_item['pk']) assert read_item['id'] == test_item['id'], "Item was not written successfully after retries" # Mock retry_utility.execute to track retries original_execute = _retry_utility.ExecuteFunction me = self.MockExecuteFunction(original_execute) _retry_utility.ExecuteFunction = me # Test without retry_write for create_item test_item_create = {"id": "2", "pk": "test", "data": "retryable create test"} me.counter = 0 # Reset counter try: container.create_item(test_item_create) pytest.fail("Expected an exception without retry_write") except exceptions.CosmosHttpResponseError: assert me.counter == 1, "Expected no retries without retry_write" # Test with retry_write for create_item me.counter = 0 # Reset counter try: container.create_item(test_item_create, retry_write=2) except exceptions.CosmosHttpResponseError as e: assert me.counter > 1, "Expected multiple retries due to simulated errors" finally: # Restore the original retry_utility.execute function _retry_utility.ExecuteFunction = original_execute # Verify the item was created read_item_create = container.read_item(item=test_item_create['id'], partition_key=test_item_create['pk']) assert read_item_create['id'] == test_item_create['id'], "Item was not created successfully after retries" # Mock retry_utility.execute to track retries original_execute = _retry_utility.ExecuteFunction me = self.MockExecuteFunction(original_execute) _retry_utility.ExecuteFunction = me # Test without retry_write for patch_item test_patch_operations = [ {"op": "add", "path": "/data", "value": "patched retryable write test"} ] try: container.patch_item(item=test_item['id'], partition_key=test_item['pk'], patch_operations=test_patch_operations) pytest.fail("Expected an exception without retry_write") except exceptions.CosmosHttpResponseError: assert me.counter == 1, "Expected no retries without retry_write" # Test with retry_write for patch_item me.counter = 0 # Reset counter try: container.patch_item(item=test_item['id'], partition_key=test_item['pk'], patch_operations=test_patch_operations, retry_write=2) except exceptions.CosmosHttpResponseError as e: assert me.counter > 1, "Expected multiple retries due to simulated errors" finally: # Restore the original retry_utility.execute function _retry_utility.ExecuteFunction = original_execute # Verify the item was patched patched_item = container.read_item(item=test_item['id'], partition_key=test_item['pk']) assert patched_item['data'] == "patched retryable write test", "Item was not patched successfully after retries" # Verify original execution function is used to upsert an item test_item = {"id": "1", "pk": "test", "data": "retryable write test"} container.upsert_item(test_item) # Mock retry_utility.execute to track retries original_execute = _retry_utility.ExecuteFunction me = self.MockExecuteFunction(original_execute) _retry_utility.ExecuteFunction = me # Verify delete_item does not retry me.counter = 0 # Reset counter try: container.delete_item(item=test_item['id'], partition_key=test_item['pk']) pytest.fail("Expected an exception without retry_write") except exceptions.CosmosHttpResponseError: assert me.counter == 1, "Expected no retries for delete_item" # Verify the item was not deleted read_item = container.read_item(item=test_item['id'], partition_key=test_item['pk']) assert read_item['id'] == test_item['id'], "Item was unexpectedly deleted" me.counter = 0 # Reset counter try: container.delete_item(item=test_item['id'], partition_key=test_item['pk'], retry_write=2) except exceptions.CosmosHttpResponseError: assert me.counter > 1, "Expected multiple retries due to simulated errors" finally: # Restore the original retry_utility.execute function _retry_utility.ExecuteFunction = original_execute try: container.read_item(item=test_item['id'], partition_key=test_item['pk']) pytest.fail("Expected an exception for item that should have been deleted") except exceptions.CosmosHttpResponseError as e: assert e.status_code == 404, "Expected item to be deleted but it was not" # Test replace item with retry_write test_item_replace = {"id": "3", "pk": "test", "data": "original data"} container.create_item(test_item_replace) # Mock retry_utility.execute to track retries original_execute = _retry_utility.ExecuteFunction me = self.MockExecuteFunction(original_execute) _retry_utility.ExecuteFunction = me # Test without retry_write for replace_item try: container.replace_item(item=test_item_replace['id'], body={"id": "3", "pk": "test", "data": "updated data"}) pytest.fail("Expected an exception without retry_write") except exceptions.CosmosHttpResponseError: assert me.counter == 1, "Expected no retries without retry_write" # Test with retry_write for replace_item me.counter = 0 # Reset counter try: container.replace_item(item=test_item_replace['id'], body={"id": "3", "pk": "test", "data": "updated data"}, retry_write=2) except exceptions.CosmosHttpResponseError as e: assert me.counter > 1, "Expected multiple retries due to simulated errors" finally: # Restore the original retry_utility.execute function _retry_utility.ExecuteFunction = original_execute # Verify the item was replaced replaced_item = container.read_item(item=test_item_replace['id'], partition_key=test_item_replace['pk']) assert replaced_item['data'] == "updated data", "Item was not replaced successfully after retries" def test_retryable_writes_client_level(self, setup): """Test retryable writes for a container with retry_write set at the client level.""" # Create a client with retry_write enabled client_with_retry = CosmosClient(self.host, credential=self.master_key, retry_write=2) container = client_with_retry.get_database_client(self.TEST_DATABASE_ID).get_container_client(setup['container'].id) # Mock retry_utility.execute to track retries original_execute = _retry_utility.ExecuteFunction me = self.MockExecuteFunction(original_execute) _retry_utility.ExecuteFunction = me # Test upsert_item test_item = {"id": "4", "pk": "test", "data": "retryable write test"} try: container.upsert_item(test_item) except exceptions.CosmosHttpResponseError as e: assert me.counter > 1, "Expected multiple retries due to simulated errors" # Test create_item me.counter = 0 # Reset counter test_item_create = {"id": "5", "pk": "test", "data": "retryable create test"} try: container.create_item(test_item_create) except exceptions.CosmosHttpResponseError as e: assert me.counter > 1, "Expected multiple retries due to simulated errors" # Test patch_item me.counter = 0 # Reset counter test_patch_operations = [ {"op": "add", "path": "/data", "value": "patched retryable write test"} ] try: container.patch_item(item=test_item['id'], partition_key=test_item['pk'], patch_operations=test_patch_operations) pytest.fail("Expected an exception without retries") except exceptions.CosmosHttpResponseError: assert me.counter == 1, "Expected no retries for patch_item with client_retry_write" finally: # Restore the original retry_utility.execute function _retry_utility.ExecuteFunction = original_execute # Verify original execution function is used to upsert an item test_item = {"id": "4", "pk": "test", "data": "retryable write test"} container.upsert_item(test_item) # Mock retry_utility.execute to track retries original_execute = _retry_utility.ExecuteFunction me = self.MockExecuteFunction(original_execute) _retry_utility.ExecuteFunction = me # Verify delete_item does not retry me.counter = 0 # Reset counter try: container.delete_item(item=test_item['id'], partition_key=test_item['pk']) except exceptions.CosmosHttpResponseError: assert me.counter > 1, "Expected multiple retries due to simulated errors" finally: # Restore the original retry_utility.execute function _retry_utility.ExecuteFunction = original_execute # Verify the item was deleted try: container.read_item(item=test_item['id'], partition_key=test_item['pk']) pytest.fail("Expected an exception for item that should have been deleted") except exceptions.CosmosHttpResponseError as e: assert e.status_code == 404, "Expected item to be deleted but it was not" # Test client level retry for service response error mf = self.MockExecuteServiceResponseException(Exception) try: # Reset the function to reset the counter _retry_utility.ExecuteFunction = mf container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())}) pytest.fail("Exception was not raised.") except ServiceResponseError: assert mf.counter == 2 finally: _retry_utility.ExecuteFunction = original_execute # Test with replace item test_item_replace = {"id": "8", "pk": "test", "data": "original data"} container.create_item(test_item_replace) # Mock retry_utility.execute to track retries original_execute = _retry_utility.ExecuteFunction me = self.MockExecuteFunction(original_execute) _retry_utility.ExecuteFunction = me # Test replace item, writes will retry by default with client level retry_write try: container.replace_item(item=test_item_replace['id'], body={"id": "8", "pk": "test", "data": "updated data"}) except exceptions.CosmosHttpResponseError as e: assert me.counter > 1, "Expected multiple retries due to simulated errors" finally: # Restore the original retry_utility.execute function _retry_utility.ExecuteFunction = original_execute # Verify the item was replaced replaced_item = container.read_item(item=test_item_replace['id'], partition_key=test_item_replace['pk']) assert replaced_item['data'] == "updated data", "Item was not replaced successfully after retries" def test_retryable_writes_hpk_request_level(self, setup): """Test retryable writes for a container with hierarchical partition keys.""" container = setup["container_hpk"] # Mock retry_utility.execute to track retries original_execute = _retry_utility.ExecuteFunction me = self.MockExecuteFunction(original_execute) _retry_utility.ExecuteFunction = me # Test without retry_write for upsert_item test_item = {"id": "1", "state": "CA", "city": "San Francisco", "data": "retryable write test"} try: container.upsert_item(test_item) pytest.fail("Expected an exception without retry_write") except exceptions.CosmosHttpResponseError: assert me.counter == 1, "Expected no retries without retry_write" # Test with retry_write for upsert_item me.counter = 0 # Reset counter try: container.upsert_item(test_item, retry_write=2) except exceptions.CosmosHttpResponseError as e: assert me.counter > 1, "Expected multiple retries due to simulated errors" finally: # Restore the original retry_utility.execute function _retry_utility.ExecuteFunction = original_execute # Verify the item was written read_item = container.read_item(item=test_item['id'], partition_key=[test_item['state'], test_item['city']]) assert read_item['id'] == test_item['id'], "Item was not written successfully after retries" # Mock retry_utility.execute to track retries original_execute = _retry_utility.ExecuteFunction me = self.MockExecuteFunction(original_execute) _retry_utility.ExecuteFunction = me # Test without retry_write for create_item test_item_create = {"id": "2", "state": "NY", "city": "New York", "data": "retryable create test"} try: container.create_item(test_item_create) pytest.fail("Expected an exception without retry_write") except exceptions.CosmosHttpResponseError: assert me.counter == 1, "Expected no retries without retry_write" # Test with retry_write for create_item me.counter = 0 # Reset counter try: container.create_item(test_item_create, retry_write=2) except exceptions.CosmosHttpResponseError as e: assert me.counter > 1, "Expected multiple retries due to simulated errors" finally: # Restore the original retry_utility.execute function _retry_utility.ExecuteFunction = original_execute # Verify the item was created read_item_create = container.read_item(item=test_item_create['id'], partition_key=[test_item_create['state'], test_item_create['city']]) assert read_item_create['id'] == test_item_create['id'], "Item was not created successfully after retries" # Mock retry_utility.execute to track retries original_execute = _retry_utility.ExecuteFunction me = self.MockExecuteFunction(original_execute) _retry_utility.ExecuteFunction = me # Test without retry_write for patch_item test_patch_operations_hpk = [ {"op": "add", "path": "/data", "value": "patched retryable write test"} ] try: container.patch_item(item=test_item['id'], partition_key=[test_item['state'], test_item['city']], patch_operations=test_patch_operations_hpk) pytest.fail("Expected an exception without retry_write") except exceptions.CosmosHttpResponseError: assert me.counter == 1, "Expected no retries without retry_write" # Test with retry_write for patch_item me.counter = 0 # Reset counter try: container.patch_item(item=test_item['id'], partition_key=[test_item['state'], test_item['city']], patch_operations=test_patch_operations_hpk, retry_write=2) except exceptions.CosmosHttpResponseError as e: assert me.counter > 1, "Expected multiple retries due to simulated errors" finally: # Restore the original retry_utility.execute function _retry_utility.ExecuteFunction = original_execute # Verify the item was patched patched_item_hpk = container.read_item(item=test_item['id'], partition_key=[test_item['state'], test_item['city']]) assert patched_item_hpk[ 'data'] == "patched retryable write test", "Item was not patched successfully after retries" # Verify original execution function is used to upsert an item test_item = {"id": "1", "state": "CA", "city": "San Francisco", "data": "retryable write test"} container.upsert_item(test_item) # Mock retry_utility.execute to track retries original_execute = _retry_utility.ExecuteFunction me = self.MockExecuteFunction(original_execute) _retry_utility.ExecuteFunction = me # Verify delete_item does not retry me.counter = 0 # Reset counter try: container.delete_item(item=test_item['id'], partition_key=[test_item['state'], test_item['city']]) except exceptions.CosmosHttpResponseError: assert me.counter == 1, "Expected no retries for delete_item" me.counter = 0 # Reset counter try: container.delete_item(item=test_item['id'], partition_key=[test_item['state'], test_item['city']], retry_write=2) except exceptions.CosmosHttpResponseError: assert me.counter > 1, "Expected multiple retries due to simulated errors" finally: # Restore the original retry_utility.execute function _retry_utility.ExecuteFunction = original_execute # Verify the item was deleted try: container.read_item(item=test_item['id'], partition_key=[test_item['state'], test_item['city']]) pytest.fail("Expected an exception for item that should have been deleted") except exceptions.CosmosHttpResponseError as e: assert e.status_code == 404, "Expected item to be deleted but it was not" # Test with replace item test_item_replace = {"id": "6", "state": "CA", "city": "San Francisco", "data": "original data"} container.create_item(test_item_replace) # Mock retry_utility.execute to track retries original_execute = _retry_utility.ExecuteFunction me = self.MockExecuteFunction(original_execute) _retry_utility.ExecuteFunction = me # Test replace item without retry_write try: container.replace_item(item=test_item_replace['id'], body={"id": "6", "state": "CA", "city": "San Francisco", "data": "updated data"}) pytest.fail("Expected an exception without retry_write") except exceptions.CosmosHttpResponseError: assert me.counter == 1, "Expected no retries without retry_write" # Test replace item with retry_write me.counter = 0 # Reset counter try: container.replace_item(item=test_item_replace['id'], body={"id": "6", "state": "CA", "city": "San Francisco", "data": "updated data"}, retry_write=2) except exceptions.CosmosHttpResponseError as e: assert me.counter > 1, "Expected multiple retries due to simulated errors" finally: # Restore the original retry_utility.execute function _retry_utility.ExecuteFunction = original_execute # Verify the item was replaced replaced_item_hpk = container.read_item(item=test_item_replace['id'], partition_key=[test_item_replace['state'], test_item_replace['city']]) assert replaced_item_hpk['data'] == "updated data", "Item was not replaced successfully after retries" def test_retryable_writes_hpk_client_level(self, setup): """Test retryable writes for a container with hierarchical partition keys and retry_write set at the client level.""" # Create a client with retry_write enabled client_with_retry = CosmosClient(self.host, credential=self.master_key, retry_write=2) container = client_with_retry.get_database_client(self.TEST_DATABASE_ID).get_container_client( setup['container_hpk'].id) # Mock retry_utility.execute to track retries original_execute = _retry_utility.ExecuteFunction me = self.MockExecuteFunction(original_execute) _retry_utility.ExecuteFunction = me # Test upsert_item test_item = {"id": "6", "state": "CA", "city": "San Francisco", "data": "retryable write test"} try: container.upsert_item(test_item) except exceptions.CosmosHttpResponseError as e: assert me.counter > 1, "Expected multiple retries due to simulated errors" # Test create_item me.counter = 0 # Reset counter test_item_create = {"id": "7", "state": "NY", "city": "New York", "data": "retryable create test"} try: container.create_item(test_item_create) except exceptions.CosmosHttpResponseError as e: assert me.counter > 1, "Expected multiple retries due to simulated errors" # Test patch_item me.counter = 0 # Reset counter test_patch_operations_hpk = [ {"op": "add", "path": "/data", "value": "patched retryable write test"} ] try: container.patch_item(item=test_item['id'], partition_key=[test_item['state'], test_item['city']], patch_operations=test_patch_operations_hpk) pytest.fail("Expected an exception without retries") except exceptions.CosmosHttpResponseError: assert me.counter == 1, "Expected no retries for patch_item with client_retry_write" finally: # Restore the original retry_utility.execute function _retry_utility.ExecuteFunction = original_execute # Verify original execution function is used to upsert an item test_item = {"id": "6", "state": "CA", "city": "San Francisco", "data": "retryable write test"} container.upsert_item(test_item) # Mock retry_utility.execute to track retries original_execute = _retry_utility.ExecuteFunction me = self.MockExecuteFunction(original_execute) _retry_utility.ExecuteFunction = me # Verify delete_item does not retry me.counter = 0 # Reset counter try: container.delete_item(item=test_item['id'], partition_key=[test_item['state'], test_item['city']]) except exceptions.CosmosHttpResponseError: assert me.counter > 1, "Expected multiple retries due to simulated errors" finally: # Restore the original retry_utility.execute function _retry_utility.ExecuteFunction = original_execute # Verify the item was deleted try: container.read_item(item=test_item['id'], partition_key=[test_item['state'], test_item['city']]) pytest.fail("Expected an exception for item that should have been deleted") except exceptions.CosmosHttpResponseError as e: assert e.status_code == 404, "Expected item to be deleted but it was not" # Test client level retry for replace item test_item_replace = {"id": "8", "state": "CA", "city": "San Francisco", "data": "original data"} container.create_item(test_item_replace) # Mock retry_utility.execute to track retries original_execute = _retry_utility.ExecuteFunction me = self.MockExecuteFunction(original_execute) _retry_utility.ExecuteFunction = me # Test replace item, writes will retry by default with client level retry_write try: container.replace_item(item=test_item_replace['id'], body={"id": "8", "state": "CA", "city": "San Francisco", "data": "updated data"}) except exceptions.CosmosHttpResponseError as e: assert me.counter > 1, "Expected multiple retries due to simulated errors" finally: # Restore the original retry_utility.execute function _retry_utility.ExecuteFunction = original_execute # Verify the item was replaced replaced_item_hpk = container.read_item(item=test_item_replace['id'], partition_key=[test_item_replace['state'], test_item_replace['city']]) assert replaced_item_hpk['data'] == "updated data", "Item was not replaced successfully after retries" @pytest.mark.parametrize("injected_error", [FaultInjectionTransport.error_request_timeout(), FaultInjectionTransport.error_internal_server_error(), FaultInjectionTransport.error_service_response()]) def test_retryable_writes_mrw_request_level(self, injected_error): # Verify cross regional retries when using retryable writes at request level with multiple write locations. # Need to test for 408s, 500s, and ServiceResponseError. first_region_uri: str = test_config.TestConfig.local_host.replace("localhost", "127.0.0.1") second_region_uri: str = test_config.TestConfig.local_host custom_transport = FaultInjectionTransport() # Inject topology transformation that would make Emulator look like a multiple write region account is_get_account_predicate: Callable[[HttpRequest], bool] = lambda \ r: FaultInjectionTransport.predicate_is_database_account_call(r) emulator_as_multi_write_region_account_transformation = \ lambda r, inner: FaultInjectionTransport.transform_topology_mwr( first_region_name="First Region", second_region_name="Second Region", inner=inner) custom_transport.add_response_transformation( is_get_account_predicate, emulator_as_multi_write_region_account_transformation) # Inject rule to simulate regional outage in "First Region" to test retry happens in next region is_request_to_first_region: Callable[[HttpRequest], bool] = lambda \ r: FaultInjectionTransport.predicate_targets_region(r, first_region_uri) and \ FaultInjectionTransport.predicate_is_document_operation(r) custom_transport.add_fault( is_request_to_first_region, lambda r: injected_error) id_value: str = str(uuid.uuid4()) document_definition = {'id': id_value, 'pk': id_value, 'name': 'sample document', 'key': 'value'} initialized_objects = self.setup_method_with_custom_transport( custom_transport, preferred_locations=["First Region", "Second Region"], multiple_write_locations=True ) container: ContainerProxy = initialized_objects["col"] create_document = container.create_item(body=document_definition, retry_write=2) request = create_document.get_response_headers()["_request"] assert request.url.startswith(second_region_uri) upsert_document = container.upsert_item(body=document_definition, retry_write=2) request = upsert_document.get_response_headers()["_request"] assert request.url.startswith(second_region_uri) replace_document = container.replace_item(item=document_definition['id'], body=document_definition, retry_write=2) request = replace_document.get_response_headers()["_request"] assert request.url.startswith(second_region_uri) operations = [{"op": "add", "path": "/favorite_color", "value": "red"},] patch_document = container.patch_item(item=document_definition['id'], partition_key=document_definition['pk'], patch_operations=operations, retry_write=2) request = patch_document.get_response_headers()["_request"] assert request.url.startswith(second_region_uri) @pytest.mark.parametrize("injected_error", [FaultInjectionTransport.error_request_timeout(), FaultInjectionTransport.error_internal_server_error(), FaultInjectionTransport.error_service_response()]) def test_retryable_writes_mrw_client_level(self, injected_error): # Verify cross regional retries when using retryable writes at client level with multiple write locations. # Need to test for 408s, 500s, and ServiceResponseError. first_region_uri: str = test_config.TestConfig.local_host.replace("localhost", "127.0.0.1") second_region_uri: str = test_config.TestConfig.local_host custom_transport = FaultInjectionTransport() # Inject topology transformation that would make Emulator look like a multiple write region account is_get_account_predicate: Callable[[HttpRequest], bool] = lambda \ r: FaultInjectionTransport.predicate_is_database_account_call(r) emulator_as_multi_write_region_account_transformation = \ lambda r, inner: FaultInjectionTransport.transform_topology_mwr( first_region_name="First Region", second_region_name="Second Region", inner=inner) custom_transport.add_response_transformation( is_get_account_predicate, emulator_as_multi_write_region_account_transformation) # Inject rule to simulate regional outage in "First Region" to test retry happens in next region is_request_to_first_region: Callable[[HttpRequest], bool] = lambda \ r: FaultInjectionTransport.predicate_targets_region(r, first_region_uri) and \ FaultInjectionTransport.predicate_is_document_operation(r) custom_transport.add_fault( is_request_to_first_region, lambda r: injected_error) id_value: str = str(uuid.uuid4()) document_definition = {'id': id_value, 'pk': id_value, 'name': 'sample document', 'key': 'value'} initialized_objects = self.setup_method_with_custom_transport( custom_transport, preferred_locations=["First Region", "Second Region"], multiple_write_locations=True, retry_write=2 ) container: ContainerProxy = initialized_objects["col"] create_document = container.create_item(body=document_definition) request = create_document.get_response_headers()["_request"] assert request.url.startswith(second_region_uri) upsert_document = container.upsert_item(body=document_definition) request = upsert_document.get_response_headers()["_request"] assert request.url.startswith(second_region_uri) replace_document = container.replace_item(item=document_definition['id'], body=document_definition) request = replace_document.get_response_headers()["_request"] assert request.url.startswith(second_region_uri) # Patch operation should not retry operations = [{"op": "add", "path": "/favorite_color", "value": "red"},] try: container.patch_item(item=document_definition['id'], partition_key=document_definition['pk'], patch_operations=operations) except (ServiceResponseError, exceptions.CosmosHttpResponseError) as e: if isinstance(e, exceptions.CosmosHttpResponseError): assert e.status_code in [408, 500], "Expected a retryable error for patch_item" class MockExecuteFunction(object): def __init__(self, org_func): self.org_func = org_func self.counter = 0 def __call__(self, func, *args, **kwargs): self.counter = self.counter + 1 if self.counter < 2: # Simulate a retryable error on the second call raise exceptions.CosmosHttpResponseError( status_code=408, message="Simulated timeout error." ) else: return self.org_func(func, *args, **kwargs) class MockExecuteServiceResponseException(object): def __init__(self, err_type): self.err_type = err_type self.counter = 0 def __call__(self, func, *args, **kwargs): self.counter = self.counter + 1 exception = ServiceResponseError("mock exception") exception.exc_type = self.err_type raise exception if __name__ == "__main__": unittest.main()