# Licensed under the MIT License.
# Copyright (c) Microsoft Corporation. All rights reserved.

import asyncio
import unittest

import pytest
import pytest_asyncio
import uuid
from azure.core.exceptions import ServiceResponseError
from azure.core.pipeline.transport._aiohttp import AioHttpTransport
from azure.core.rest import HttpRequest
from azure.cosmos.aio import CosmosClient, _retry_utility_async
from azure.cosmos import exceptions, ContainerProxy, DatabaseProxy
from azure.cosmos.partition_key import PartitionKey
from _fault_injection_transport_async import FaultInjectionTransportAsync
from test_fault_injection_transport_async import TestFaultInjectionTransportAsync
from typing import Callable
import test_config  # Import test_config for configuration details

HPK_CONTAINER = "TestContainerHPK"

@pytest_asyncio.fixture()
async 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 = await test_client.create_database_if_not_exists(id=config.TEST_DATABASE_ID)
    created_container = await database.create_container_if_not_exists(
        id=config.TEST_SINGLE_PARTITION_CONTAINER_ID,
        partition_key=PartitionKey(path="/pk")
    )
    container_hpk = await 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.asyncio
@pytest.mark.usefixtures("setup")
class TestRetryableWritesAsync:
    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
    async def setup_method_with_custom_transport(
            custom_transport: AioHttpTransport,
            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, enable_diagnostics_logging=True, **kwargs)
        await client.__aenter__()
        db: DatabaseProxy = await client.create_database_if_not_exists(database_id)
        container: ContainerProxy = await db.create_container_if_not_exists(container_id, PartitionKey(path="/pk"))
        return {"client": client, "db": db, "col": container}

    async def test_retryable_writes_request_level_async(self, setup):
        container = setup['container']

        # Mock retry_utility.execute to track retries
        original_execute = _retry_utility_async.ExecuteFunctionAsync
        me = self.MockExecuteFunction(original_execute)
        _retry_utility_async.ExecuteFunctionAsync = me

        # Test without retry_write for upsert_item
        test_item = {"id": "1", "pk": "test", "data": "retryable write test"}
        try:
            await 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:
            await 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_async.ExecuteFunctionAsync = original_execute

        # Verify the item was written
        read_item = await 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_async.ExecuteFunctionAsync
        me = self.MockExecuteFunction(original_execute)
        _retry_utility_async.ExecuteFunctionAsync = me

        # Test without retry_write for create_item
        test_item_create = {"id": "2", "pk": "test", "data": "retryable create test"}
        try:
            await 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:
            await 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_async.ExecuteFunctionAsync = original_execute

        # Verify the item was created
        read_item_create = await 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_async.ExecuteFunctionAsync
        me = self.MockExecuteFunction(original_execute)
        _retry_utility_async.ExecuteFunctionAsync = me

        # Test without retry_write for patch_item
        test_patch_operations = [
            {"op": "add", "path": "/data", "value": "patched retryable write test"}
        ]
        try:
            await container.patch_item(item=test_item_create['id'], partition_key=test_item_create['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:
            await container.patch_item(item=test_item_create['id'], partition_key=test_item_create['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_async.ExecuteFunctionAsync = original_execute

        # Verify the item was patched
        patched_item = await container.read_item(item=test_item_create['id'], partition_key=test_item_create['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": "5", "pk": "test", "data": "retryable write test"}
        await container.upsert_item(test_item)

        # Mock retry_utility.execute to track retries
        original_execute = _retry_utility_async.ExecuteFunctionAsync
        me = self.MockExecuteFunction(original_execute)
        _retry_utility_async.ExecuteFunctionAsync = me

        # Verify delete_item does not retry
        me.counter = 0  # Reset counter
        try:
            await 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"
        finally:
            # Restore the original retry_utility.execute function
            _retry_utility_async.ExecuteFunctionAsync = original_execute

        # Verify the item was not deleted
        read_item = await container.read_item(item=test_item['id'], partition_key=test_item['pk'])
        assert read_item['id'] == test_item['id'], "Item was unexpectedly deleted"

        # Mock retry_utility.execute to track retries
        original_execute = _retry_utility_async.ExecuteFunctionAsync
        me = self.MockExecuteFunction(original_execute)
        _retry_utility_async.ExecuteFunctionAsync = me

        me.counter = 0  # Reset counter
        try:
            await 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_async.ExecuteFunctionAsync = original_execute

        try:
            await 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"

        # Now we test it out with a replace request without retry write enabled - which should not retry

        # First create an item to replace
        test_item_replace = {"id": str(uuid.uuid4()), "pk": "test", "data": "replace test"}
        await container.create_item(test_item_replace)

        # Mock execute function
        original_execute = _retry_utility_async.ExecuteFunctionAsync
        me = self.MockExecuteFunction(original_execute)
        _retry_utility_async.ExecuteFunctionAsync = me
        me.counter = 0  # Reset counter
        try:
            await container.replace_item(item=test_item_replace['id'], body={"id": test_item_replace['id'],
                                                                              "pk": "test",
                                                                              "data": "original data"})
            pytest.fail("Expected an exception without retry_write")
        except exceptions.CosmosHttpResponseError:
            assert me.counter == 1, "Expected no retries for replace_item without retry_write"

        # Test with retry_write for replace_item
        me.counter = 0  # Reset counter
        try:
            await container.replace_item(item=test_item_replace['id'], body={"id": test_item_replace['id'],
                                                                              "pk": "test",
                                                                              "data": "replaced 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_async.ExecuteFunctionAsync = original_execute

        # Verify the item was replaced
        read_item_replace = await container.read_item(item=test_item_replace['id'],
                                                      partition_key=test_item_replace['pk'])
        assert read_item_replace['data'] == "replaced data", "Item was not replaced successfully after retries"


    async def test_retryable_writes_client_level_write(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)
        await client_with_retry.__aenter__()

        try:
            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_async.ExecuteFunctionAsync
            me = self.MockExecuteFunction(original_execute)
            _retry_utility_async.ExecuteFunctionAsync = me

            # Test upsert_item
            test_item = {"id": "4", "pk": "test", "data": "retryable write test"}
            try:
                await 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:
                await 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:
                await 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_async.ExecuteFunction = original_execute

            # Verify original execution function is used to upsert an item
            test_item = {"id": "4", "pk": "test", "data": "retryable write test"}
            await container.upsert_item(test_item)

            # Mock retry_utility.execute to track retries
            original_execute = _retry_utility_async.ExecuteFunctionAsync
            me = self.MockExecuteFunction(original_execute)
            _retry_utility_async.ExecuteFunction = me

            # Verify delete_item does not retry
            me.counter = 0  # Reset counter
            try:
                await 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_async.ExecuteFunctionAsync = original_execute

            # Verify the item was deleted
            try:
                await 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 on client level with retry_write enabled
            # First create an item to replace
            test_item_replace = {"id": str(uuid.uuid4()), "pk": "test", "data": "original data"}
            await container.create_item(test_item_replace)
            # Mock execute function
            original_execute = _retry_utility_async.ExecuteFunctionAsync
            me = self.MockExecuteFunction(original_execute)
            _retry_utility_async.ExecuteFunctionAsync = me
            me.counter = 0  # Reset counter
            try:
                await container.replace_item(item=test_item_replace['id'], body={"id": test_item_replace['id'],
                                                                                  "pk": "test",
                                                                                  "data": "replaced data"})
            except exceptions.CosmosHttpResponseError:
                assert me.counter > 1, "Expected multiple retries for replace_item with retry_write enabled"
            finally:
                # Restore the original retry_utility.execute function
                _retry_utility_async.ExecuteFunctionAsync = original_execute
            # Verify the item was replaced
            read_item_replace = await container.read_item(item=test_item_replace['id'],
                                                          partition_key=test_item_replace['pk'])
            assert read_item_replace['data'] == "replaced data", "Item was not replaced successfully after retries"
        finally:
            if client_with_retry:
                await client_with_retry.close()

    async def test_retryable_writes_hpk_request_level_async(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_async.ExecuteFunctionAsync
        me = self.MockExecuteFunction(original_execute)
        _retry_utility_async.ExecuteFunctionAsync = me

        # Test without retry_write for upsert_item
        test_item = {"id": "1", "state": "CA", "city": "San Francisco", "data": "retryable write test"}
        try:
            await 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:
            await 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_async.ExecuteFunctionAsync = original_execute

        # Verify the item was written
        read_item = await 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_async.ExecuteFunctionAsync
        me = self.MockExecuteFunction(original_execute)
        _retry_utility_async.ExecuteFunctionAsync = me

        # Test without retry_write for create_item
        test_item_create = {"id": "2", "state": "NY", "city": "New York", "data": "retryable create test"}
        try:
            await 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:
            await 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_async.ExecuteFunctionAsync = original_execute

        # Verify the item was created
        read_item_create = await 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_async.ExecuteFunctionAsync
        me = self.MockExecuteFunction(original_execute)
        _retry_utility_async.ExecuteFunctionAsync = me

        test_patch_operations = [
            {"op": "add", "path": "/year", "value": "2025"}
        ]
        # Test without retry_write for patch_item
        try:
            await container.patch_item(item=test_item['id'],
                                       partition_key=[test_item['state'], test_item['city']],
                                       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"

        # Now properly replace with retry_write
        me.counter = 0  # Reset counter
        await container.patch_item(item=test_item['id'],
                                    partition_key=[test_item['state'], test_item['city']],
                                    patch_operations=test_patch_operations,
                                    retry_write=2)
        assert me.counter == 2

        # Verify the item was patched
        patched_item = await container.read_item(item=test_item['id'],
                                                        partition_key=[test_item['state'],
                                                                       test_item['city']])
        assert patched_item['year'] == "2025", "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"}
        await container.upsert_item(test_item)

        # Mock retry_utility.execute to track retries
        original_execute = _retry_utility_async.ExecuteFunctionAsync
        me = self.MockExecuteFunction(original_execute)
        _retry_utility_async.ExecuteFunctionAsync = me

        # Verify delete_item does not retry
        me.counter = 0  # Reset counter
        try:
            await 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:
            await 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_async.ExecuteFunctionAsync = original_execute

        # Verify the item was deleted
        try:
            await 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"

        # Now we test it out with a replace request without retry write enabled - which should not retry
        # First create an item to replace
        test_item_replace = {"id": str(uuid.uuid4()), "state": "TX", "city": "Austin", "data": "replace test"}
        await container.create_item(test_item_replace)
        # Mock execute function
        original_execute = _retry_utility_async.ExecuteFunctionAsync
        me = self.MockExecuteFunction(original_execute)
        _retry_utility_async.ExecuteFunctionAsync = me
        me.counter = 0  # Reset counter
        try:
            await container.replace_item(item=test_item_replace['id'], body={"id": test_item_replace['id'],
                                                                              "state": "TX",
                                                                              "city": "Austin",
                                                                              "data": "original data"})
            pytest.fail("Expected an exception without retry_write")
        except exceptions.CosmosHttpResponseError:
            assert me.counter == 1, "Expected no retries for replace_item without retry_write"
        # Test with retry_write for replace_item
        me.counter = 0  # Reset counter
        try:
            await container.replace_item(item=test_item_replace['id'], body={"id": test_item_replace['id'],
                                                                              "state": "TX",
                                                                              "city": "Austin",
                                                                              "data": "replaced 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_async.ExecuteFunctionAsync = original_execute

        # Verify the item was replaced
        read_item_replace = await container.read_item(item=test_item_replace['id'],
                                                        partition_key=[test_item_replace['state'],
                                                                         test_item_replace['city']])
        assert read_item_replace['data'] == "replaced data", "Item was not replaced successfully after retries"

    async def test_retryable_writes_hpk_client_level_async(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)
        await client_with_retry.__aenter__()
        try:
            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_async.ExecuteFunctionAsync
            me = self.MockExecuteFunction(original_execute)
            _retry_utility_async.ExecuteFunctionAsync = me

            # Test upsert_item
            test_item = {"id": "6", "state": "CA", "city": "San Francisco", "data": "retryable write test"}
            try:
                await 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:
                await 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:
                await 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_async.ExecuteFunction = original_execute

            test_item = {"id": "6", "state": "CA", "city": "San Francisco", "data": "retryable write test"}
            await container.upsert_item(test_item)

            # Mock retry_utility.execute to track retries
            original_execute = _retry_utility_async.ExecuteFunctionAsync
            me = self.MockExecuteFunction(original_execute)
            _retry_utility_async.ExecuteFunctionAsync = me

            # Verify delete_item does not retry
            me.counter = 0  # Reset counter
            try:
                await 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_async.ExecuteFunctionAsync = original_execute

            # Verify the item was deleted
            try:
                await 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"

            # Now we try a replace request with retry write enabled - which should retry once

            # Create an item to replace
            test_item_replace = {"id": str(uuid.uuid4()), "state": "TX", "city": "Austin", "data": "original data"}
            await container.create_item(test_item_replace)
            # Mock execute function
            original_execute = _retry_utility_async.ExecuteFunctionAsync
            me = self.MockExecuteFunction(original_execute)
            _retry_utility_async.ExecuteFunctionAsync = me
            me.counter = 0  # Reset counter
            try:
                await container.replace_item(item=test_item_replace['id'], body={"id": test_item_replace['id'],
                                                                                  "state": "TX",
                                                                                  "city": "Austin",
                                                                                  "data": "replaced data"})
            except exceptions.CosmosHttpResponseError:
                assert me.counter > 1, "Expected multiple retries for replace_item with retry_write enabled"
            finally:
                # Restore the original retry_utility.execute function
                _retry_utility_async.ExecuteFunctionAsync = original_execute
            # Verify the item was replaced
            read_item_replace = await container.read_item(item=test_item_replace['id'],
                                                          partition_key=[test_item_replace['state'],
                                                                         test_item_replace['city']])
            assert read_item_replace['data'] == "replaced data", "Item was not replaced successfully after retries"

        finally:
            await client_with_retry.close()

    @pytest.mark.parametrize("injected_error", [FaultInjectionTransportAsync.error_request_timeout(),
                             FaultInjectionTransportAsync.error_internal_server_error(),
                             FaultInjectionTransportAsync.error_service_response()])
    async def test_retryable_writes_mrw_request_level_async(self, injected_error):
        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 = FaultInjectionTransportAsync()

        # Inject topology transformation that would make Emulator look like a multiple write region account
        # account with two read regions
        is_get_account_predicate: Callable[[HttpRequest], bool] = lambda \
            r: FaultInjectionTransportAsync.predicate_is_database_account_call(r)
        emulator_as_multi_write_region_account_transformation = \
            lambda r, inner: FaultInjectionTransportAsync.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"
        is_request_to_first_region: Callable[[HttpRequest], bool] = lambda \
                r: FaultInjectionTransportAsync.predicate_targets_region(r, first_region_uri) and \
                   FaultInjectionTransportAsync.predicate_is_document_operation(r)

        custom_transport.add_fault(
            is_request_to_first_region,
            lambda r: asyncio.create_task(FaultInjectionTransportAsync.error_region_down()))

        id_value: str = str(uuid.uuid4())
        document_definition = {'id': id_value,
                               'pk': id_value,
                               'name': 'sample document',
                               'key': 'value'}

        initialized_objects = await self.setup_method_with_custom_transport(
            custom_transport,
            preferred_locations=["First Region", "Second Region"],
            multiple_write_locations=True
        )
        try:
            container: ContainerProxy = initialized_objects["col"]

            create_document = await 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 = await 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 = await 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 = await 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)

        finally:
            await TestFaultInjectionTransportAsync.cleanup_method(initialized_objects)

    @pytest.mark.parametrize("injected_error", [FaultInjectionTransportAsync.error_request_timeout(),
                             FaultInjectionTransportAsync.error_internal_server_error(),
                             FaultInjectionTransportAsync.error_service_response()])
    async def test_retryable_writes_mrw_client_level_async(self, injected_error):
        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 = FaultInjectionTransportAsync()

        # Inject topology transformation that would make Emulator look like a multiple write region account
        is_get_account_predicate: Callable[[HttpRequest], bool] = lambda \
            r: FaultInjectionTransportAsync.predicate_is_database_account_call(r)
        emulator_as_multi_write_region_account_transformation = \
            lambda r, inner: FaultInjectionTransportAsync.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"
        is_request_to_first_region: Callable[[HttpRequest], bool] = lambda \
                r: FaultInjectionTransportAsync.predicate_targets_region(r, first_region_uri) and \
                   FaultInjectionTransportAsync.predicate_is_document_operation(r)

        custom_transport.add_fault(
            is_request_to_first_region,
            lambda r: asyncio.create_task(FaultInjectionTransportAsync.error_region_down()))

        id_value: str = str(uuid.uuid4())
        document_definition = {'id': id_value,
                               'pk': id_value,
                               'name': 'sample document',
                               'key': 'value'}

        initialized_objects = await self.setup_method_with_custom_transport(
            custom_transport,
            preferred_locations=["First Region", "Second Region"],
            multiple_write_locations=True,
            retry_write=2
        )
        try:
            container: ContainerProxy = initialized_objects["col"]

            create_document = await container.create_item(body=document_definition)
            request = create_document.get_response_headers()["_request"]
            assert request.url.startswith(second_region_uri)

            upsert_document = await container.upsert_item(body=document_definition)
            request = upsert_document.get_response_headers()["_request"]
            assert request.url.startswith(second_region_uri)

            replace_document = await 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:
                await 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"

        finally:
            await TestFaultInjectionTransportAsync.cleanup_method(initialized_objects)

    class MockExecuteFunction(object):
        def __init__(self, org_func):
            self.org_func = org_func
            self.counter = 0

        async 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 await self.org_func(func, *args, **kwargs)

    class MockExecuteServiceResponseException(object):
        def __init__(self, err_type, inner_exception):
            self.err_type = err_type
            self.inner_exception = inner_exception
            self.counter = 0

        def __call__(self, func, *args, **kwargs):
            self.counter = self.counter + 1
            exception = ServiceResponseError("mock exception")
            exception.exc_type = self.err_type
            exception.inner_exception = self.inner_exception
            raise exception

if __name__ == '__main__':
    unittest.main()