# The MIT License (MIT)
# Copyright (c) Microsoft Corporation. All rights reserved.

import unittest
import uuid
from unittest.mock import patch
from urllib.parse import urlparse

import pytest
from aiohttp.client_exceptions import (ClientConnectionError, ClientConnectionResetError,
                                       ClientOSError, ServerConnectionError)
from azure.core.exceptions import ServiceRequestError, ServiceResponseError

import test_config
from azure.cosmos import DatabaseAccount, _location_cache
from azure.cosmos._location_cache import RegionalRoutingContext
from azure.cosmos.aio import CosmosClient, _retry_utility_async, _global_endpoint_manager_async
from azure.cosmos.exceptions import CosmosHttpResponseError


@pytest.mark.cosmosEmulator
class TestServiceRetryPoliciesAsync(unittest.IsolatedAsyncioTestCase):
    host = test_config.TestConfig.host
    masterKey = test_config.TestConfig.masterKey
    connectionPolicy = test_config.TestConfig.connectionPolicy
    TEST_DATABASE_ID = test_config.TestConfig.TEST_DATABASE_ID
    TEST_CONTAINER_ID = test_config.TestConfig.TEST_SINGLE_PARTITION_CONTAINER_ID
    REGION1 = "West US"
    REGION2 = "East US"
    REGION3 = "West US 2"

    @classmethod
    def setUpClass(cls):
        if (cls.masterKey == '[YOUR_KEY_HERE]' or
                cls.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.")

    async def asyncSetUp(self):
        self.client = CosmosClient(self.host, self.masterKey)
        self.created_database = self.client.get_database_client(self.TEST_DATABASE_ID)
        self.created_container = self.created_database.get_container_client(self.TEST_CONTAINER_ID)

    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

    @classmethod
    def _uses_localhost_endpoint(cls):
        parsed = urlparse(cls.host)
        return parsed.hostname in ("localhost", "127.0.0.1")

    @classmethod
    def _make_regional_endpoint(cls, region):
        """Return a per-region locational endpoint (e.g. ``acct-westus...``).

        Each region must have its own endpoint URL. If every region shares
        the default endpoint, ``LocationCache.is_default_endpoint_regional``
        becomes True and the next ``update_location_cache`` call clears
        ``effective_preferred_locations`` and shrinks
        ``read_regional_routing_contexts`` to a single fallback, turning
        multi-region tests into one-shot tests. Cancelling the background
        refresh task is not enough on its own: foreground refreshes can fire
        from the pkranges side-call that the IgnoreQuery mocks let through.
        """
        return RegionalRoutingContext(
            _location_cache.LocationCache.GetLocationalEndpoint(cls.host, region)
        )

    def _setup_read_regions(self, location_cache, regions):
        """Populate the read side of the location cache with N distinct regions.

        Mirrors the production initialization flow: set the raw inputs
        (account locations, locational-endpoint map, preferred locations),
        clear any unavailability state carried over from a previous
        assertion, then call ``update_location_cache()`` so the derived
        dicts (``_read_locations_by_normalized``, the reverse endpoint→region
        map, and ``read_regional_routing_contexts``) are recomputed from a
        consistent snapshot. Direct attribute assignment alone leaves stale
        derived state behind, which silently inflates the retry budget the
        next assertion observes.
        """
        if self._uses_localhost_endpoint():
            shared_context = RegionalRoutingContext(self.host)
            location_cache.account_read_locations = list(regions)
            location_cache.account_read_regional_routing_contexts_by_location = {
                r: shared_context for r in regions
            }
            location_cache.account_locations_by_read_endpoints = {self.host: regions[0]}
            location_cache.effective_preferred_locations = list(regions)
            location_cache.read_regional_routing_contexts = [shared_context for _ in regions]
            location_cache.location_unavailability_info_by_endpoint = {}
            return

        endpoints_by_region = {r: self._make_regional_endpoint(r) for r in regions}
        location_cache.account_read_locations = list(regions)
        location_cache.account_read_regional_routing_contexts_by_location = endpoints_by_region
        # Reverse map (endpoint URL -> region name). The retry policy uses
        # this to translate ``location_endpoint_to_route`` back to a region
        # when marking endpoints unavailable; if it is stale the wrong
        # region gets marked and the retry budget can drift.
        location_cache.account_locations_by_read_endpoints = {
            ctx.get_primary(): name for name, ctx in endpoints_by_region.items()
        }
        location_cache.effective_preferred_locations = list(regions)
        # Each assertion in this test reuses the same location cache; reset
        # unavailability so a region marked unavailable in the previous
        # 3-region step does not silently shrink (or extend) the next step's
        # effective routing list.
        location_cache.location_unavailability_info_by_endpoint = {}
        # Recompute derived state from the raw inputs above so the helper's
        # output matches what the production initialization path would
        # produce for the same topology.
        location_cache.update_location_cache()

    def _setup_write_regions(self, location_cache, regions):
        """Populate the write side of the location cache with N distinct regions.

        Companion to ``_setup_read_regions`` — see that docstring for the
        rationale on clearing unavailability state and re-running
        ``update_location_cache()``.
        """
        if self._uses_localhost_endpoint():
            shared_context = RegionalRoutingContext(self.host)
            location_cache.account_write_locations = list(regions)
            location_cache.account_write_regional_routing_contexts_by_location = {
                r: shared_context for r in regions
            }
            location_cache.account_locations_by_write_endpoints = {self.host: regions[0]}
            location_cache.write_regional_routing_contexts = [shared_context for _ in regions]
            location_cache.location_unavailability_info_by_endpoint = {}
            return

        endpoints_by_region = {r: self._make_regional_endpoint(r) for r in regions}
        location_cache.account_write_locations = list(regions)
        location_cache.account_write_regional_routing_contexts_by_location = endpoints_by_region
        location_cache.account_locations_by_write_endpoints = {
            ctx.get_primary(): name for name, ctx in endpoints_by_region.items()
        }
        location_cache.location_unavailability_info_by_endpoint = {}
        location_cache.update_location_cache()

    def _setup_shared_write_endpoint_regions(self, location_cache, regions):
        """Set write routing contexts to one shared endpoint for all regions.

        ``test_service_response_errors_async`` intentionally validates the path
        where two in-region write entries point to the same endpoint URL
        (``self.host``). Keep this as derived-state setup only (no
        ``update_location_cache()`` call), matching the historical test shape.
        """
        shared_context = RegionalRoutingContext(self.host)
        location_cache.account_write_locations = list(regions)
        location_cache.write_regional_routing_contexts = [shared_context for _ in 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:
            db = mock_client.get_database_client(self.TEST_DATABASE_ID)
            container = db.get_container_client(self.TEST_CONTAINER_ID)

            created_item = await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())})

            # 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
            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)
            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

            # Now we test with a query operation
            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()
            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

            # Now we change the location cache to have only 1 preferred read region
            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):
                with pytest.raises(ServiceRequestError):
                    await container.read_item(created_item['id'], created_item['pk'])
                assert mf.counter == expected_counter

            # Now we try it out with a write request
            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):
                with pytest.raises(ServiceRequestError):
                    await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())})
                assert mf.counter == expected_counter

    async def test_service_response_retry_policy_async(self):
        # For ServiceResponseErrors, we only do cross region retries on read requests or on ClientConnectionErrors
        # We also only do retries within the ConnectionRetryPolicy for the cases above
        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)

            created_item = await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())})

            # 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
            self._setup_read_regions(original_location_cache, [self.REGION1, self.REGION2, self.REGION3])

            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

            # Now we change the location cache to have only 1 preferred read region
            self._setup_read_regions(original_location_cache, [self.REGION1])
            expected_counter = len(original_location_cache.read_regional_routing_contexts)
            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 == expected_counter

            # Now we try it out with a write request
            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):
                    # 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

            # If we do a write request with a ClientConnectionError,
            # we will do cross-region retries like with read requests
            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):
                    await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())})
                assert mf.counter == 2

            # Now we try it out with a write request with retry write enabled - which should retry once
            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

    async def test_service_request_connection_retry_policy_async(self):
        # Mock the client retry policy to see the same-region retries that happen there
        exception = ServiceRequestError("mock exception")
        exception.exc_type = Exception
        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) as mock_client:
            db = mock_client.get_database_client(self.TEST_DATABASE_ID)
            container = db.get_container_client(self.TEST_CONTAINER_ID)

            # We retry ServiceRequestExceptions 3 times in the to the same endpoint before raising the exception
            # regardless of operation type
            try:
                await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())})
                pytest.fail("Exception was not raised.")
            except ServiceRequestError:
                assert connection_retry_policy.counter == 3

            try:
                await container.read_item("some_id", "some_pk")
                pytest.fail("Exception was not raised.")
            except ServiceRequestError:
                assert connection_retry_policy.counter == 3

    async def test_service_response_connection_retry_policy_async(self):
        # Mock the client retry policy to see the same-region retries that happen there
        exception = ServiceResponseError("mock exception")
        exception.exc_type = Exception
        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) as mock_client:
            db = mock_client.get_database_client(self.TEST_DATABASE_ID)
            container = db.get_container_client(self.TEST_CONTAINER_ID)

            # We retry ServiceResponseExceptions 3 times in the to the same endpoint before raising the exception
            # for read operations, but 0 times for write requests
            try:
                await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())})
                pytest.fail("Exception was not raised.")
            except ServiceResponseError:
                assert connection_retry_policy.counter == 0

            try:
                await container.read_item("some_id", "some_pk")
                pytest.fail("Exception was not raised.")
            except ServiceResponseError:
                assert connection_retry_policy.counter == 3

    async def test_service_response_errors_async(self):
        # Test for errors that are subclasses of ClientConnectionError for write requests
        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)
            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
            self._setup_read_regions(original_location_cache, [self.REGION1, self.REGION2, self.REGION3])
            # For writes, keep two entries that share one endpoint URL. This
            # test validates the shared-endpoint unavailability path.
            self._setup_shared_write_endpoint_regions(original_location_cache, [self.REGION1, self.REGION2])

            # 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

            # Now we test the base ClientConnectionError to see in-region retry
            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')

            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
            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')

            # Reset the location cache's unavailable endpoints in order to try the same with other exceptions
            original_location_cache.location_unavailability_info_by_endpoint = {}
            self._setup_shared_write_endpoint_regions(original_location_cache, [self.REGION1, self.REGION2])

            # Now we test ClientConnectionResetError, the subclass of ClientConnectionError
            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')

            # Reset the location cache's unavailable endpoints in order to try the same with other exceptions
            original_location_cache.location_unavailability_info_by_endpoint = {}
            self._setup_shared_write_endpoint_regions(original_location_cache, [self.REGION1, self.REGION2])

            # Now we test ServerConnectionError, the subclass of ClientConnectionError
            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

            # Reset the location cache's unavailable endpoints in order to try the same with other exceptions
            original_location_cache.location_unavailability_info_by_endpoint = {}
            self._setup_shared_write_endpoint_regions(original_location_cache, [self.REGION1, self.REGION2])

            # Now we test ClientOSError, the subclass of ClientConnectionError
            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

    async def test_global_endpoint_manager_retry_async(self):
        # For this test we mock both the ConnectionRetryPolicy and the GetDatabaseAccountStub
        # - ConnectionRetryPolicy allows us to raise Service exceptions only for chosen requests and track endpoints used
        # - GetDatabaseAccountStub allows us to receive any number of endpoints for that call independent of account used
        exception = ServiceRequestError("mock exception")
        exception.exc_type = Exception
        connection_policy = self.connectionPolicy
        connection_retry_policy = test_config.MockConnectionRetryPolicyAsync(resource_type="docs", error=exception)
        connection_policy.ConnectionRetryConfiguration = connection_retry_policy
        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

                # 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

    class MockExecuteServiceRequestException(object):
        def __init__(self):
            self.counter = 0

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

    class MockExecuteServiceRequestExceptionIgnoreQuery(object):
        def __init__(self, original_execute_function):
            self.counter = 0
            self.original_execute_function = original_execute_function

        def __call__(self, func, *args, **kwargs):
            if len(args) > 1:
                request_obj = args[1]
                if not (hasattr(request_obj, "resource_type") and hasattr(request_obj, "operation_type")):
                    return self.original_execute_function(func, *args, **kwargs)
                if ((request_obj.resource_type == "docs" and request_obj.operation_type == "Query") or
                        (request_obj.resource_type == "pkranges" and request_obj.operation_type == "ReadFeed")):
                    # Ignore query requests, As an additional ReadFeed might occur during a regular Read operation
                    return self.original_execute_function(func, *args, **kwargs)
                self.counter = self.counter + 1
                exception = ServiceRequestError("mock exception")
                exception.exc_type = Exception
                raise exception
            return self.original_execute_function(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

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

        def __call__(self, func, *args, **kwargs):
            if len(args) > 1:
                request_obj = args[1]
                if not (hasattr(request_obj, "resource_type") and hasattr(request_obj, "operation_type")):
                    return self.original_execute_function(func, *args, **kwargs)
                if ((request_obj.resource_type == "docs" and request_obj.operation_type == "Query") or
                        (request_obj.resource_type == "pkranges" and request_obj.operation_type == "ReadFeed")):
                    # Ignore query requests, As an additional ReadFeed might occur during a regular Read operation
                    return self.original_execute_function(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
            return self.original_execute_function(func, *args, **kwargs)

    async def MockGetDatabaseAccountStub(self, endpoint):
        read_regions = ["West US", "East US"]
        read_locations = []
        for loc in read_regions:
            read_locations.append({'databaseAccountEndpoint': endpoint, 'name': loc})
        write_regions = ["West US"]
        write_locations = []
        for loc in write_regions:
            locational_endpoint = self.host.replace("localhost", "127.0.0.1")
            write_locations.append({'databaseAccountEndpoint': locational_endpoint, 'name': loc})
        multi_write = False

        db_acc = DatabaseAccount()
        db_acc.DatabasesLink = "/dbs/"
        db_acc.MediaLink = "/media/"
        db_acc._ReadableLocations = read_locations
        db_acc._WritableLocations = write_locations
        db_acc._EnableMultipleWritableLocations = multi_write
        db_acc.ConsistencyPolicy = {"defaultConsistencyLevel": "Session"}
        return db_acc
