# ------------------------------------
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
# ------------------------------------
from devtools_testutils.perfstress_tests import PerfStressTest

from azure.identity import ClientSecretCredential, TokenCachePersistenceOptions
from azure.identity.aio import ClientSecretCredential as AsyncClientSecretCredential

try:
    from dotenv import load_dotenv

    load_dotenv()
except ImportError:
    pass


class PersistentCacheRead(PerfStressTest):
    def __init__(self, arguments):
        super().__init__(arguments)

        client_id = self.get_from_env("AZURE_CLIENT_ID")
        tenant_id = self.get_from_env("AZURE_TENANT_ID")
        secret = self.get_from_env("AZURE_CLIENT_SECRET")
        cache_options = TokenCachePersistenceOptions(allow_unencrypted_storage=True)
        self.credential = ClientSecretCredential(tenant_id, client_id, secret, cache_persistence_options=cache_options)
        self.async_credential = AsyncClientSecretCredential(
            tenant_id, client_id, secret, cache_persistence_options=cache_options
        )
        self.scope = "https://vault.azure.net/.default"

    async def global_setup(self):
        """Cache an access token"""

        await super().global_setup()
        self.credential.get_token(self.scope)
        await self.async_credential.get_token(self.scope)

    def run_sync(self):
        self.credential.get_token(self.scope)

    async def run_async(self):
        await self.async_credential.get_token(self.scope)

    async def close(self):
        await self.async_credential.close()
        await super().close()
