# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
import pytest

from devtools_testutils import AzureTestCase, ResponseCallback


from azure.core.exceptions import (
    HttpResponseError,
    ResourceExistsError,
    AzureError,
)
from azure.core.pipeline.policies import RetryMode
from azure.core.pipeline.transport import(
    AioHttpTransport
)

from azure.data.tables.aio import TableServiceClient

from _shared.asynctestcase import AsyncTableTestCase

from async_preparers import tables_decorator_async


class RetryAioHttpTransport(AioHttpTransport):
    """Transport to test retry"""
    def __init__(self, *args, **kwargs):
        super(RetryAioHttpTransport, self).__init__(*args, **kwargs)
        self.count = 0

    async def send(self, request, **kwargs):
        self.count += 1
        response = await super(RetryAioHttpTransport, self).send(request, **kwargs)
        return response


# --Test Class -----------------------------------------------------------------
class StorageRetryTest(AzureTestCase, AsyncTableTestCase):

    async def _set_up(self, tables_storage_account_name, tables_primary_storage_account_key, url='table', default_table=True, **kwargs):
        self.table_name = self.get_resource_name('uttable')
        self.ts = TableServiceClient(
            self.account_url(tables_storage_account_name, url),
            credential=tables_primary_storage_account_key,
            **kwargs
        )
        self.table = self.ts.get_table_client(self.table_name)
        if self.is_live and default_table:
            try:
                await self.ts.create_table(self.table_name)
            except ResourceExistsError:
                pass

        self.query_tables = []

    # --Test Cases --------------------------------------------
    @tables_decorator_async
    async def test_retry_on_server_error_async(self, tables_storage_account_name, tables_primary_storage_account_key):
        await self._set_up(tables_storage_account_name, tables_primary_storage_account_key, default_table=False)
        try:
            callback = ResponseCallback(status=201, new_status=500).override_status

            new_table_name = self.get_resource_name('uttable')
            # The initial create will return 201, but we overwrite it with 500 and retry.
            # The retry will then get a 409 conflict.
            with pytest.raises(ResourceExistsError):
                await self.ts.create_table(new_table_name, raw_response_hook=callback)
        finally:
            await self.ts.delete_table(new_table_name)
            await self._tear_down()

    @tables_decorator_async
    async def test_retry_on_timeout_async(self, tables_storage_account_name, tables_primary_storage_account_key):
        await self._set_up(
            tables_storage_account_name,
            tables_primary_storage_account_key,
            default_table=False,
            retry_mode=RetryMode.Exponential,
            retry_backoff_factor=1)

        callback = ResponseCallback(status=200, new_status=408).override_first_status

        try:
            # The initial create will return 201, but we overwrite it with 408 and retry.
            # The retry will then succeed.
            await self.ts.get_service_properties(raw_response_hook=callback)
        finally:
            await self._tear_down()

    @pytest.mark.live_test_only
    @tables_decorator_async
    async def test_retry_on_socket_timeout_async(self, tables_storage_account_name, tables_primary_storage_account_key):
        retry_transport = RetryAioHttpTransport(connection_timeout=11, read_timeout=0.000000000001)
        await self._set_up(
            tables_storage_account_name,
            tables_primary_storage_account_key,
            retry_mode=RetryMode.Fixed,
            retry_backoff_factor=1,
            transport=retry_transport,
            default_table=False)

        with pytest.raises(AzureError) as error:
            await self.ts.get_service_properties()

        # 3 retries + 1 original == 4
        assert retry_transport.count == 4
        # This call should succeed on the server side, but fail on the client side due to socket timeout
        self.assertTrue('Timeout on reading' in str(error.value), 'Expected socket timeout but got different exception.')

    @tables_decorator_async
    async def test_no_retry_async(self, tables_storage_account_name, tables_primary_storage_account_key):
        await self._set_up(tables_storage_account_name, tables_primary_storage_account_key, retry_total=0, default_table=False)

        new_table_name = self.get_resource_name('uttable')

        # Force the create call to error with a 500
        callback = ResponseCallback(status=201, new_status=500).override_status

        try:
            with pytest.raises(HttpResponseError) as error:
                await self.ts.create_table(new_table_name, raw_response_hook=callback)
            assert error.value.response.status_code == 500
            assert error.value.reason == 'Created'

        finally:
            await self.ts.delete_table(new_table_name)
            await self._tear_down()
# ------------------------------------------------------------------------------
