# pylint: disable=line-too-long,useless-suppression
# ------------------------------------
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
# ------------------------------------
import random
import re
import functools
import json
from typing import Optional, Any, Dict
from azure.ai.projects.models import (
    Connection,
    ConnectionType,
    CustomCredential,
    CredentialType,
    ApiKeyCredentials,
    Deployment,
    DeploymentType,
    ModelDeployment,
    Index,
    IndexType,
    AzureAISearchIndex,
    DatasetVersion,
    DatasetType,
    DatasetCredential,
    ItemResource,
    ItemType,
    ResponsesMessageRole,
    ItemContentType,
)
from azure.ai.projects.models._models import AgentObject, AgentVersionObject
from devtools_testutils import AzureRecordedTestCase, EnvironmentVariableLoader, is_live_and_not_recording
from azure.ai.projects import AIProjectClient as AIProjectClient
from azure.ai.projects.aio import AIProjectClient as AsyncAIProjectClient

# temporary requirements for recorded_by_proxy_httpx and recorded_by_proxy_async_httpx
import logging
import urllib.parse as url_parse

try:
    import httpx
except ImportError:
    httpx = None

from azure.core.exceptions import ResourceNotFoundError
from azure.core.pipeline.policies import ContentDecodePolicy

from devtools_testutils import is_live_internal, is_live_and_not_recording, trim_kwargs_from_test_function
from devtools_testutils.proxy_testcase import (
    get_test_id,
    start_record_or_playback,
    stop_record_or_playback,
    get_proxy_netloc,
)

# Load secrets from environment variables
servicePreparer = functools.partial(
    EnvironmentVariableLoader,
    "azure_ai_projects_tests",
    azure_ai_projects_tests_project_endpoint="https://sanitized-account-name.services.ai.azure.com/api/projects/sanitized-project-name",
    azure_ai_projects_tests_agents_project_endpoint="https://sanitized-account-name.services.ai.azure.com/api/projects/sanitized-project-name",
    azure_ai_projects_tests_tracing_project_endpoint="https://sanitized-account-name.services.ai.azure.com/api/projects/sanitized-project-name",
    azure_ai_projects_tests_container_app_resource_id="/subscriptions/00000000-0000-0000-0000-000000000000/resourceGroups/00000/providers/Microsoft.App/containerApps/00000",
    azure_ai_projects_tests_container_ingress_subdomain_suffix="00000",
)


class TestBase(AzureRecordedTestCase):

    test_redteams_params = {
        # cSpell:disable-next-line
        "connection_name": "naposaniwestus3",
        "connection_type": ConnectionType.AZURE_OPEN_AI,
        "model_deployment_name": "gpt-4o-mini",
    }

    test_connections_params = {
        "connection_name": "custom_keys_connection",
        "connection_type": ConnectionType.CUSTOM,
    }

    test_deployments_params = {
        "model_publisher": "Cohere",
        "model_name": "gpt-4o",
        "model_deployment_name": "DeepSeek-V3",
    }

    test_agents_params = {
        "model_deployment_name": "gpt-4o",
        "agent_name": "agent-for-python-projects-sdk-testing",
    }

    test_inference_params = {
        "connection_name_api_key_auth": "connection1",
        "connection_name_entra_id_auth": "connection2",
        "model_deployment_name": "gpt-4o",
        "aoai_api_version": "2025-04-01-preview",
    }

    test_indexes_params = {
        "index_name": f"test-index-name",
        "index_version": "1",
        "ai_search_connection_name": "my-ai-search-connection",
        "ai_search_index_name": "my-ai-search-index",
    }

    test_datasets_params = {
        "dataset_name_1": f"test-dataset-name-{random.randint(0, 99999):05d}",
        "dataset_name_2": f"test-dataset-name-{random.randint(0, 99999):05d}",
        "dataset_name_3": f"test-dataset-name-{random.randint(0, 99999):05d}",
        "dataset_name_4": f"test-dataset-name-{random.randint(0, 99999):05d}",
        "dataset_version": 1,
        "connection_name": "balapvbyostoragecanary",
    }

    test_files_params = {
        "test_file_name": "test_file.jsonl",
        "file_purpose": "fine-tune",
    }

    test_finetuning_params = {
        "sft": {
            "openai": {
                "model_name": "gpt-4.1",
                "deployment": {
                    "deployment_name": "gpt-4-1-fine-tuned-test",
                    "pre_finetuned_model": "ft:gpt-4.1:azure-ai-test::ABCD1234",
                },
            },
            "oss": {"model_name": "Ministral-3B"},
            "training_file_name": "sft_training_set.jsonl",
            "validation_file_name": "sft_validation_set.jsonl",
        },
        "dpo": {
            "openai": {"model_name": "gpt-4o-mini"},
            "training_file_name": "dpo_training_set.jsonl",
            "validation_file_name": "dpo_validation_set.jsonl",
        },
        "rft": {
            "openai": {"model_name": "o4-mini"},
            "training_file_name": "rft_training_set.jsonl",
            "validation_file_name": "rft_validation_set.jsonl",
        },
        "n_epochs": 1,
        "batch_size": 1,
        "learning_rate_multiplier": 1.0,
    }

    # Regular expression describing the pattern of an Application Insights connection string.
    REGEX_APPINSIGHTS_CONNECTION_STRING = re.compile(
        r"^InstrumentationKey=[0-9a-fA-F-]{36};IngestionEndpoint=https://.+.applicationinsights.azure.com/;LiveEndpoint=https://.+.monitor.azure.com/;ApplicationId=[0-9a-fA-F-]{36}$"
    )

    # helper function: create projects client using environment variables
    def create_client(self, *, operation_group: Optional[str] = None, **kwargs) -> AIProjectClient:
        # fetch environment variables
        project_endpoint_env_variable = (
            f"azure_ai_projects_tests_{operation_group}_project_endpoint"
            if operation_group
            else "azure_ai_projects_tests_project_endpoint"
        )
        endpoint = kwargs.pop(project_endpoint_env_variable)
        credential = self.get_credential(AIProjectClient, is_async=False)

        # create and return client
        client = AIProjectClient(
            endpoint=endpoint,
            credential=credential,
        )

        return client

    # helper function: create async projects client using environment variables
    def create_async_client(self, *, operation_group: Optional[str] = None, **kwargs) -> AsyncAIProjectClient:
        # fetch environment variables
        project_endpoint_env_variable = (
            f"azure_ai_projects_tests_{operation_group}_project_endpoint"
            if operation_group
            else "azure_ai_projects_tests_project_endpoint"
        )
        endpoint = kwargs.pop(project_endpoint_env_variable)
        credential = self.get_credential(AsyncAIProjectClient, is_async=True)

        # create and return async client
        client = AsyncAIProjectClient(
            endpoint=endpoint,
            credential=credential,
        )

        return client

    @staticmethod
    def assert_equal_or_not_none(actual, expected=None):
        assert actual is not None
        if expected is not None:
            assert actual == expected

    # Checks that a given dictionary has at least one non-empty (non-whitespace) string key-value pair.
    @classmethod
    def is_valid_dict(cls, d: dict[str, str]) -> bool:
        return bool(d) and all(
            isinstance(k, str) and isinstance(v, str) and k.strip() and v.strip() for k, v in d.items()
        )

    @classmethod
    def validate_connection(
        cls,
        connection: Connection,
        include_credentials: bool,
        *,
        expected_connection_type: Optional[ConnectionType] = None,
        expected_connection_name: Optional[str] = None,
        expected_authentication_type: Optional[CredentialType] = None,
        expected_is_default: Optional[bool] = None,
    ):
        assert connection.id is not None

        TestBase.assert_equal_or_not_none(connection.name, expected_connection_name)
        TestBase.assert_equal_or_not_none(connection.type, expected_connection_type)
        TestBase.assert_equal_or_not_none(connection.credentials.type, expected_authentication_type)

        if expected_is_default is not None:
            assert connection.is_default == expected_is_default

        if isinstance(connection.credentials, ApiKeyCredentials):
            assert connection.credentials.type == CredentialType.API_KEY
            if include_credentials:
                assert connection.credentials.api_key is not None
        elif isinstance(connection.credentials, CustomCredential):
            assert connection.credentials.type == CredentialType.CUSTOM
            if include_credentials:
                assert TestBase.is_valid_dict(connection.credentials.credential_keys)

    @classmethod
    def validate_red_team_response(
        cls, response, expected_attack_strategies: int = -1, expected_risk_categories: int = -1
    ):
        """Assert basic red team scan response properties."""
        assert response is not None
        assert hasattr(response, "name")
        assert hasattr(response, "display_name")
        assert hasattr(response, "status")
        assert hasattr(response, "attack_strategies")
        assert hasattr(response, "risk_categories")
        assert hasattr(response, "target")
        assert hasattr(response, "properties")

        # Validate attack strategies and risk categories
        if expected_attack_strategies != -1:
            assert len(response.attack_strategies) == expected_attack_strategies
        if expected_risk_categories != -1:
            assert len(response.risk_categories) == expected_risk_categories
        assert response.status is not None
        cls._assert_azure_ml_properties(response)

    @classmethod
    def _assert_azure_ml_properties(cls, response):
        """Assert Azure ML specific properties are present and valid."""
        properties = response.properties
        assert properties is not None, "Red team scan properties should not be None"

        required_properties = [
            "runType",
            "redteaming",
            "_azureml.evaluation_run",
            "_azureml.evaluate_artifacts",
            "AiStudioEvaluationUri",
        ]

        for prop in required_properties:
            assert prop in properties, f"Missing required property: {prop}"

        # Validate specific property values
        assert properties["runType"] == "eval_run"
        assert properties["_azureml.evaluation_run"] == "evaluation.service"
        assert "instance_results.json" in properties["_azureml.evaluate_artifacts"]
        assert properties["redteaming"] == "asr"

        # Validate AI Studio URI format
        ai_studio_uri = properties["AiStudioEvaluationUri"]
        assert ai_studio_uri.startswith("https://ai.azure.com/resource/build/redteaming/")
        assert "wsid=" in ai_studio_uri
        assert "tid=" in ai_studio_uri

    @classmethod
    def validate_deployment(
        cls,
        deployment: Deployment,
        *,
        expected_model_name: Optional[str] = None,
        expected_model_deployment_name: Optional[str] = None,
        expected_model_publisher: Optional[str] = None,
    ):
        assert type(deployment) == ModelDeployment
        assert deployment.type == DeploymentType.MODEL_DEPLOYMENT
        assert deployment.model_version is not None
        # Comment out the below, since I see that `Cohere-embed-v3-english` has an empty capabilities dict.
        # assert TestBase.is_valid_dict(deployment.capabilities)
        assert bool(deployment.sku)  # Check none-empty

        TestBase.assert_equal_or_not_none(deployment.model_name, expected_model_name)
        TestBase.assert_equal_or_not_none(deployment.name, expected_model_deployment_name)
        TestBase.assert_equal_or_not_none(deployment.model_publisher, expected_model_publisher)

    @classmethod
    def validate_index(
        cls,
        index: Index,
        *,
        expected_index_type: Optional[IndexType] = None,
        expected_index_name: Optional[str] = None,
        expected_index_version: Optional[str] = None,
        expected_ai_search_connection_name: Optional[str] = None,
        expected_ai_search_index_name: Optional[str] = None,
    ):

        TestBase.assert_equal_or_not_none(index.name, expected_index_name)
        TestBase.assert_equal_or_not_none(index.version, expected_index_version)

        if expected_index_type == IndexType.AZURE_SEARCH:
            assert type(index) == AzureAISearchIndex
            assert index.type == IndexType.AZURE_SEARCH
            TestBase.assert_equal_or_not_none(index.connection_name, expected_ai_search_connection_name)
            TestBase.assert_equal_or_not_none(index.index_name, expected_ai_search_index_name)

    @classmethod
    def validate_dataset(
        cls,
        dataset: DatasetVersion,
        *,
        expected_dataset_type: Optional[DatasetType] = None,
        expected_dataset_name: Optional[str] = None,
        expected_dataset_version: Optional[str] = None,
        expected_connection_name: Optional[str] = None,
    ):
        assert dataset.data_uri is not None

        if expected_dataset_type:
            assert dataset.type == expected_dataset_type
        else:
            assert dataset.type == DatasetType.URI_FILE or dataset.type == DatasetType.URI_FOLDER

        TestBase.assert_equal_or_not_none(dataset.name, expected_dataset_name)
        TestBase.assert_equal_or_not_none(dataset.version, expected_dataset_version)
        if expected_connection_name:
            assert dataset.connection_name == expected_connection_name

    @classmethod
    def validate_dataset_credential(cls, dataset_credential: DatasetCredential):

        assert dataset_credential.blob_reference is not None
        assert dataset_credential.blob_reference.blob_uri
        assert dataset_credential.blob_reference.storage_account_arm_id

        assert dataset_credential.blob_reference.credential is not None
        assert (
            dataset_credential.blob_reference.credential.type == "SAS"
        )  # Why is this not of type CredentialType.SAS as defined for Connections?
        assert dataset_credential.blob_reference.credential.sas_uri

    @staticmethod
    def _validate_conversation(
        conversation: Any, *, expected_id: Optional[str] = None, expected_metadata: Optional[Dict[str, Any]] = None
    ) -> None:
        assert conversation.id
        assert conversation.created_at
        if expected_id is not None:
            assert conversation.id == expected_id
        if expected_metadata is not None:
            assert conversation.metadata == expected_metadata
        print(f"Conversation validated (id: {conversation.id})")

    def _validate_agent_version(
        self, agent: AgentVersionObject, expected_name: Optional[str] = None, expected_version: Optional[str] = None
    ) -> None:
        assert agent is not None
        assert isinstance(agent, AgentVersionObject)
        assert agent.id is not None
        if expected_name:
            assert agent.name == expected_name
        if expected_version:
            assert agent.version == expected_version
        print(f"Agent version validated (id: {agent.id}, name: {agent.name}, version: {agent.version})")

    def _validate_agent(
        self, agent: AgentObject, expected_name: Optional[str] = None, expected_latest_version: Optional[str] = None
    ) -> None:
        assert agent is not None
        assert isinstance(agent, AgentObject)
        assert agent.id is not None
        if expected_name:
            assert agent.name == expected_name
        if expected_latest_version:
            assert agent.versions.latest.version == expected_latest_version
        print(f"Agent validated (id: {agent.id}, name: {agent.name}, latest version: {agent.versions.latest.version})")

    def _validate_conversation_item(
        self,
        item: ItemResource,
        *,
        expected_type: Optional[ItemType] = None,
        expected_id: Optional[str] = None,
        expected_role: Optional[ResponsesMessageRole] = None,
        expected_content_type: Optional[ItemContentType] = None,
        expected_content_text: Optional[str] = None,
    ) -> None:
        assert item

        # From ItemResource:
        if expected_type:
            assert item.type == expected_type
        else:
            assert item.type
        if expected_id:
            assert item.id == expected_id
        else:
            assert item.id

        # From ResponsesMessageItemResource:
        if expected_type == ItemType.MESSAGE:
            assert item.status == "completed"
            if expected_role:
                assert item.role == expected_role
            else:
                assert item.role

            # From ResponsesAssistantMessageItemResource, ResponsesDeveloperMessageItemResource, ResponsesSystemMessageItemResource, ResponsesUserMessageItemResource:
            assert len(item.content) == 1

            # From ItemContent:
            if expected_content_type:
                assert item.content[0].type == expected_content_type
            if expected_content_text:
                assert item.content[0].text == expected_content_text
        print(
            f"Conversation item validated (id: {item.id}, type: {item.type}, role: {item.role if item.type == ItemType.MESSAGE else 'N/A'})"
        )

    @classmethod
    def validate_file(
        cls,
        file_obj,
        *,
        expected_file_id: Optional[str] = None,
        expected_filename: Optional[str] = None,
        expected_purpose: Optional[str] = None,
    ):
        assert file_obj is not None
        assert file_obj.id is not None
        assert file_obj.bytes is not None
        assert file_obj.created_at is not None
        assert file_obj.filename is not None
        assert file_obj.purpose is not None

        TestBase.assert_equal_or_not_none(file_obj.id, expected_file_id)
        TestBase.assert_equal_or_not_none(file_obj.filename, expected_filename)
        TestBase.assert_equal_or_not_none(file_obj.purpose, expected_purpose)

    @classmethod
    def validate_fine_tuning_job(
        cls,
        job_obj,
        *,
        expected_job_id: Optional[str] = None,
        expected_model: Optional[str] = None,
        expected_status: Optional[str] = None,
    ):
        assert job_obj is not None
        assert job_obj.id is not None
        assert job_obj.model is not None
        assert job_obj.created_at is not None
        assert job_obj.status is not None
        assert job_obj.training_file is not None

        TestBase.assert_equal_or_not_none(job_obj.id, expected_job_id)
        TestBase.assert_equal_or_not_none(job_obj.model, expected_model)
        TestBase.assert_equal_or_not_none(job_obj.status, expected_status)

    def _request_callback(self, pipeline_request) -> None:
        self.pipeline_request = pipeline_request

    @staticmethod
    def _are_json_equal(json_str1: str, json_str2: str) -> bool:
        try:
            obj1 = json.loads(json_str1)
            obj2 = json.loads(json_str2)
            return obj1 == obj2
        except json.JSONDecodeError as e:
            print(f"Invalid JSON: {e}")
            return False


# The following two decorators recorded_by_proxy_httpx and recorded_by_proxy_async_httpx will be supplanted as part of #43794.
# These are provided here temporarily to support existing tests that use httpx-based clients until they can be migrated.
def recorded_by_proxy_httpx(test_func):
    """Decorator that redirects httpx network requests to target the azure-sdk-tools test proxy.

    Use this decorator for tests that use httpx-based clients (like OpenAI SDK) instead of Azure SDK clients.
    It monkeypatches httpx.HTTPTransport.handle_request to route requests through the test proxy.

    For more details and usage examples, refer to
    https://github.com/Azure/azure-sdk-for-python/blob/main/doc/dev/tests.md#write-or-run-tests
    """
    if httpx is None:
        raise ImportError("httpx is required to use recorded_by_proxy_httpx. Install it with: pip install httpx")

    def record_wrap(*args, **kwargs):
        def transform_httpx_request(request: httpx.Request, recording_id: str) -> None:
            """Transform an httpx.Request to route through the test proxy."""
            parsed_result = url_parse.urlparse(str(request.url))

            # Store original upstream URI
            if "x-recording-upstream-base-uri" not in request.headers:
                request.headers["x-recording-upstream-base-uri"] = f"{parsed_result.scheme}://{parsed_result.netloc}"

            # Set recording headers
            request.headers["x-recording-id"] = recording_id
            request.headers["x-recording-mode"] = "record" if is_live_internal() else "playback"

            # Remove all request headers that start with `x-stainless`, since they contain CPU info, OS info, etc.
            # Those change depending on which machine the tests are run on, so we cannot have a single test recording with those.
            headers_to_remove = [key for key in request.headers.keys() if key.lower().startswith("x-stainless")]
            for header in headers_to_remove:
                del request.headers[header]

            # Rewrite URL to proxy
            updated_target = parsed_result._replace(**get_proxy_netloc()).geturl()
            request.url = httpx.URL(updated_target)

        def restore_httpx_response_url(response: httpx.Response) -> httpx.Response:
            """Restore the response's request URL to the original upstream target."""
            try:
                parsed_resp = url_parse.urlparse(str(response.request.url))
                upstream_uri_str = response.request.headers.get("x-recording-upstream-base-uri", "")
                if upstream_uri_str:
                    upstream_uri = url_parse.urlparse(upstream_uri_str)
                    original_target = parsed_resp._replace(
                        scheme=upstream_uri.scheme or parsed_resp.scheme, netloc=upstream_uri.netloc
                    ).geturl()
                    response.request.url = httpx.URL(original_target)
            except Exception:
                # Best-effort restore; don't fail the call if something goes wrong
                pass
            return response

        trimmed_kwargs = {k: v for k, v in kwargs.items()}
        trim_kwargs_from_test_function(test_func, trimmed_kwargs)

        if is_live_and_not_recording():
            return test_func(*args, **trimmed_kwargs)

        test_id = get_test_id()
        recording_id, variables = start_record_or_playback(test_id)
        original_transport_func = httpx.HTTPTransport.handle_request

        def combined_call(transport_self, request: httpx.Request) -> httpx.Response:
            transform_httpx_request(request, recording_id)
            result = original_transport_func(transport_self, request)
            return restore_httpx_response_url(result)

        httpx.HTTPTransport.handle_request = combined_call

        # Call the test function
        test_variables = None
        test_run = False
        try:
            try:
                test_variables = test_func(*args, variables=variables, **trimmed_kwargs)
                test_run = True
            except TypeError as error:
                if "unexpected keyword argument" in str(error) and "variables" in str(error):
                    logger = logging.getLogger()
                    logger.info(
                        "This test can't accept variables as input. The test method should accept `**kwargs` and/or a "
                        "`variables` parameter to make use of recorded test variables."
                    )
                else:
                    raise error
            # If the test couldn't accept `variables`, run without passing them
            if not test_run:
                test_variables = test_func(*args, **trimmed_kwargs)

        except ResourceNotFoundError as error:
            error_body = ContentDecodePolicy.deserialize_from_http_generics(error.response)
            message = error_body.get("message") or error_body.get("Message")
            error_with_message = ResourceNotFoundError(message=message, response=error.response)
            raise error_with_message from error

        finally:
            httpx.HTTPTransport.handle_request = original_transport_func
            stop_record_or_playback(test_id, recording_id, test_variables)

        return test_variables

    return record_wrap


def recorded_by_proxy_async_httpx(test_func):
    """Decorator that redirects async httpx network requests to target the azure-sdk-tools test proxy.

    Use this decorator for async tests that use httpx-based clients (like OpenAI AsyncOpenAI SDK)
    instead of Azure SDK clients. It monkeypatches httpx.AsyncHTTPTransport.handle_async_request
    to route requests through the test proxy.

    For more details and usage examples, refer to
    https://github.com/Azure/azure-sdk-for-python/blob/main/doc/dev/tests.md#write-or-run-tests
    """
    if httpx is None:
        raise ImportError("httpx is required to use recorded_by_proxy_async_httpx. Install it with: pip install httpx")

    async def record_wrap(*args, **kwargs):
        def transform_httpx_request(request: httpx.Request, recording_id: str) -> None:
            """Transform an httpx.Request to route through the test proxy."""
            parsed_result = url_parse.urlparse(str(request.url))

            # Store original upstream URI
            if "x-recording-upstream-base-uri" not in request.headers:
                request.headers["x-recording-upstream-base-uri"] = f"{parsed_result.scheme}://{parsed_result.netloc}"

            # Set recording headers
            request.headers["x-recording-id"] = recording_id
            request.headers["x-recording-mode"] = "record" if is_live_internal() else "playback"

            # Remove all request headers that start with `x-stainless`, since they contain CPU info, OS info, etc.
            # Those change depending on which machine the tests are run on, so we cannot have a single test recording with those.
            headers_to_remove = [key for key in request.headers.keys() if key.lower().startswith("x-stainless")]
            for header in headers_to_remove:
                del request.headers[header]

            # Rewrite URL to proxy
            updated_target = parsed_result._replace(**get_proxy_netloc()).geturl()
            request.url = httpx.URL(updated_target)

        def restore_httpx_response_url(response: httpx.Response) -> httpx.Response:
            """Restore the response's request URL to the original upstream target."""
            try:
                parsed_resp = url_parse.urlparse(str(response.request.url))
                upstream_uri_str = response.request.headers.get("x-recording-upstream-base-uri", "")
                if upstream_uri_str:
                    upstream_uri = url_parse.urlparse(upstream_uri_str)
                    original_target = parsed_resp._replace(
                        scheme=upstream_uri.scheme or parsed_resp.scheme, netloc=upstream_uri.netloc
                    ).geturl()
                    response.request.url = httpx.URL(original_target)
            except Exception:
                # Best-effort restore; don't fail the call if something goes wrong
                pass
            return response

        trimmed_kwargs = {k: v for k, v in kwargs.items()}
        trim_kwargs_from_test_function(test_func, trimmed_kwargs)

        if is_live_and_not_recording():
            return await test_func(*args, **trimmed_kwargs)

        test_id = get_test_id()
        recording_id, variables = start_record_or_playback(test_id)
        original_transport_func = httpx.AsyncHTTPTransport.handle_async_request

        async def combined_call(transport_self, request: httpx.Request) -> httpx.Response:
            transform_httpx_request(request, recording_id)
            result = await original_transport_func(transport_self, request)
            return restore_httpx_response_url(result)

        httpx.AsyncHTTPTransport.handle_async_request = combined_call

        # Call the test function
        test_variables = None
        test_run = False
        try:
            try:
                test_variables = await test_func(*args, variables=variables, **trimmed_kwargs)
                test_run = True
            except TypeError as error:
                if "unexpected keyword argument" in str(error) and "variables" in str(error):
                    logger = logging.getLogger()
                    logger.info(
                        "This test can't accept variables as input. The test method should accept `**kwargs` and/or a "
                        "`variables` parameter to make use of recorded test variables."
                    )
                else:
                    raise error
            # If the test couldn't accept `variables`, run without passing them
            if not test_run:
                test_variables = await test_func(*args, **trimmed_kwargs)

        except ResourceNotFoundError as error:
            error_body = ContentDecodePolicy.deserialize_from_http_generics(error.response)
            message = error_body.get("message") or error_body.get("Message")
            error_with_message = ResourceNotFoundError(message=message, response=error.response)
            raise error_with_message from error

        finally:
            httpx.AsyncHTTPTransport.handle_async_request = original_transport_func
            stop_record_or_playback(test_id, recording_id, test_variables)

        return test_variables

    return record_wrap