from unittest.mock import MagicMock, patch

from airflow_tools.cli.deploy import deploy


@patch("airflow_tools.cli.deploy.S3Connector")
@patch("airflow_tools.cli.deploy.MWAAConnector")
def test_deploy(mwaa_connector: MagicMock, s3_connector: MagicMock):
    deploy(
        environment="test",
        service_name="airflow-service",
        bucket="deployment-bucket",
        kms_key_id="01968e55-21a8-4968-8870-4c33a2a5f845",
        git_commit="0b96734d86343965094bb3268a5534bd18487210",
        requirements_file="requirements.txt",
        dags_folder="dags",
    )

    s3_connector.return_value.upload_file.assert_called_once_with(
        bucket="deployment-bucket",
        filename="requirements.txt",
        key="0b96734d86343965094bb3268a5534bd18487210/requirements.txt",
        kms_key_id="01968e55-21a8-4968-8870-4c33a2a5f845",
    )
    s3_connector.return_value.upload_dir.assert_called_once_with(
        bucket="deployment-bucket",
        folder="dags",
        key_prefix="0b96734d86343965094bb3268a5534bd18487210/dags",
        kms_key_id="01968e55-21a8-4968-8870-4c33a2a5f845",
    )
    mwaa_connector.return_value.update_environment.assert_called_once_with(
        environment_name="test-airflow-service",
        bucket="deployment-bucket",
        dag_s3_path="0b96734d86343965094bb3268a5534bd18487210/dags",
        requirements_s3_path="0b96734d86343965094bb3268a5534bd18487210/requirements.txt",
    )


@patch("airflow_tools.cli.deploy.S3Connector")
@patch("airflow_tools.cli.deploy.MWAAConnector")
def test_deploy_with_bucket_prefix(mwaa_connector: MagicMock, s3_connector: MagicMock):
    deploy(
        environment="test",
        service_name="airflow-service",
        bucket="deployment-bucket",
        kms_key_id="01968e55-21a8-4968-8870-4c33a2a5f845",
        git_commit="0b96734d86343965094bb3268a5534bd18487210",
        requirements_file="requirements.txt",
        dags_folder="dags",
        bucket_prefix="some-prefix",
    )

    s3_connector.return_value.upload_file.assert_called_once_with(
        bucket="deployment-bucket",
        filename="requirements.txt",
        key="some-prefix/0b96734d86343965094bb3268a5534bd18487210/requirements.txt",
        kms_key_id="01968e55-21a8-4968-8870-4c33a2a5f845",
    )
    s3_connector.return_value.upload_dir.assert_called_once_with(
        bucket="deployment-bucket",
        folder="dags",
        key_prefix="some-prefix/0b96734d86343965094bb3268a5534bd18487210/dags",
        kms_key_id="01968e55-21a8-4968-8870-4c33a2a5f845",
    )
    mwaa_connector.return_value.update_environment.assert_called_once_with(
        environment_name="test-airflow-service",
        bucket="deployment-bucket",
        dag_s3_path="some-prefix/0b96734d86343965094bb3268a5534bd18487210/dags",
        requirements_s3_path="some-prefix/0b96734d86343965094bb3268a5534bd18487210/requirements.txt",
    )


@patch("airflow_tools.cli.deploy.S3Connector")
@patch("airflow_tools.cli.deploy.MWAAConnector")
def test_deploy_with_bucket_prefix_trailing_slash(
    mwaa_connector: MagicMock, s3_connector: MagicMock
):
    deploy(
        environment="test",
        service_name="airflow-service",
        bucket="deployment-bucket",
        kms_key_id="01968e55-21a8-4968-8870-4c33a2a5f845",
        git_commit="0b96734d86343965094bb3268a5534bd18487210",
        requirements_file="requirements.txt",
        dags_folder="dags",
        bucket_prefix="some-prefix/",
    )

    s3_connector.return_value.upload_file.assert_called_once_with(
        bucket="deployment-bucket",
        filename="requirements.txt",
        key="some-prefix/0b96734d86343965094bb3268a5534bd18487210/requirements.txt",
        kms_key_id="01968e55-21a8-4968-8870-4c33a2a5f845",
    )
    s3_connector.return_value.upload_dir.assert_called_once_with(
        bucket="deployment-bucket",
        folder="dags",
        key_prefix="some-prefix/0b96734d86343965094bb3268a5534bd18487210/dags",
        kms_key_id="01968e55-21a8-4968-8870-4c33a2a5f845",
    )
    mwaa_connector.return_value.update_environment.assert_called_once_with(
        environment_name="test-airflow-service",
        bucket="deployment-bucket",
        dag_s3_path="some-prefix/0b96734d86343965094bb3268a5534bd18487210/dags",
        requirements_s3_path="some-prefix/0b96734d86343965094bb3268a5534bd18487210/requirements.txt",
    )


@patch("airflow_tools.cli.deploy.S3Connector")
@patch("airflow_tools.cli.deploy.MWAAConnector")
def test_rollback(mwaa_connector: MagicMock, s3_connector: MagicMock):
    deploy(
        environment="test",
        service_name="airflow-service",
        bucket="deployment-bucket",
        kms_key_id="01968e55-21a8-4968-8870-4c33a2a5f845",
        git_commit="0b96734d86343965094bb3268a5534bd18487210",
        requirements_file="requirements.txt",
        dags_folder="dags",
        rollback=True,
    )

    s3_connector.return_value.upload_file.assert_not_called()
    s3_connector.return_value.upload_dir.assert_not_called()
    mwaa_connector.return_value.update_environment.assert_called_once_with(
        environment_name="test-airflow-service",
        bucket="deployment-bucket",
        dag_s3_path="0b96734d86343965094bb3268a5534bd18487210/dags",
        requirements_s3_path="0b96734d86343965094bb3268a5534bd18487210/requirements.txt",
    )