from unittest.mock import MagicMock

import pytest

from airflow_tools.connectors.mwaa import MWAAConnector


def test_update_environment():
    """Test the update_environment method of MWAAConnector."""
    mock_client = MagicMock()
    mwaa_connector = MWAAConnector(retry_interval=0, timeout=0)
    mwaa_connector._client = mock_client

    # Mock the responses for get_environment
    mock_client.get_environment.side_effect = [
        {"Environment": {"Status": "AVAILABLE"}},
        {"Environment": {"LastUpdate": {"Status": "SUCCESS"}}},
    ]

    # Call the method
    mwaa_connector.update_environment(
        environment_name="test-environment",
        bucket="test-bucket",
        dag_s3_path="dags",
        requirements_s3_path="requirements.txt",
    )

    # Assert that the update_environment method was called with the correct parameters
    mock_client.update_environment.assert_called_once_with(
        Name="test-environment",
        SourceBucketArn="arn:aws:s3:::test-bucket",
        DagS3Path="dags",
        RequirementsS3Path="requirements.txt",
    )


def test_update_environment_timeout_on_availability():
    """
    Test the update_environment method of MWAAConnector when it times out
    on availability.
    """
    mock_client = MagicMock()
    mwaa_connector = MWAAConnector(retry_interval=0, timeout=0)
    mwaa_connector._client = mock_client

    mock_client.get_environment.return_value = {"Environment": {"Status": "PENDING"}}

    with pytest.raises(
        Exception,
        match="Timed out waiting for environment test-environment to be available.",
    ):
        mwaa_connector.update_environment(
            environment_name="test-environment",
            bucket="test-bucket",
            dag_s3_path="dags",
            requirements_s3_path="requirements.txt",
        )


def test_update_environment_timeout_on_update():
    """
    Test the update_environment method of MWAAConnector when it times out
    on update.
    """
    mock_client = MagicMock()
    mwaa_connector = MWAAConnector(retry_interval=0, timeout=0)
    mwaa_connector._client = mock_client

    mock_client.get_environment.side_effect = [
        {"Environment": {"Status": "AVAILABLE"}},
        {"Environment": {"LastUpdate": {"Status": "PENDING"}}},
    ]

    with pytest.raises(
        Exception, match="Timed out waiting for deployment to complete."
    ):
        mwaa_connector.update_environment(
            environment_name="test-environment",
            bucket="test-bucket",
            dag_s3_path="dags",
            requirements_s3_path="requirements.txt",
        )


def test_update_environment_deployment_failed():
    """Test the update_environment method of MWAAConnector when deployment fails."""
    mock_client = MagicMock()
    mwaa_connector = MWAAConnector(retry_interval=0, timeout=0)
    mwaa_connector._client = mock_client

    mock_client.get_environment.side_effect = [
        {"Environment": {"Status": "AVAILABLE"}},
        {
            "Environment": {
                "LastUpdate": {
                    "Status": "FAILED",
                    "Error": {
                        "ErrorCode": "SomeError",
                        "ErrorMessage": "Some error message",
                    },
                }
            }
        },
    ]

    with pytest.raises(
        Exception,
        match="Deployment failed. Error code: SomeError, message: Some error message",
    ):
        mwaa_connector.update_environment(
            environment_name="test-environment",
            bucket="test-bucket",
            dag_s3_path="dags",
            requirements_s3_path="requirements.txt",
        )