#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
from __future__ import annotations

from typing import Any, TypedDict
from unittest import mock

import pytest
from botocore.waiter import Waiter

from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.providers.amazon.aws.hooks.eks import ClusterStates, EksHook
from airflow.providers.amazon.aws.operators.eks import (
    EksCreateClusterOperator,
    EksCreateFargateProfileOperator,
    EksCreateNodegroupOperator,
    EksDeleteClusterOperator,
    EksDeleteFargateProfileOperator,
    EksDeleteNodegroupOperator,
    EksPodOperator,
)
from airflow.providers.amazon.aws.triggers.eks import (
    EksCreateFargateProfileTrigger,
    EksCreateNodegroupTrigger,
    EksDeleteFargateProfileTrigger,
)
from airflow.providers.cncf.kubernetes.utils.pod_manager import OnFinishAction
from airflow.providers.common.compat.sdk import TaskDeferred

from unit.amazon.aws.utils.eks_test_constants import (
    NODEROLE_ARN,
    POD_EXECUTION_ROLE_ARN,
    RESOURCES_VPC_CONFIG,
    ROLE_ARN,
    SELECTORS,
    SUBNET_IDS,
    TASK_ID,
)
from unit.amazon.aws.utils.eks_test_utils import convert_keys
from unit.amazon.aws.utils.test_template_fields import validate_template_fields
from unit.amazon.aws.utils.test_waiter import assert_expected_waiter_type

CLUSTER_NAME = "cluster1"
NODEGROUP_NAME = "nodegroup1"
FARGATE_PROFILE_NAME = "fargate_profile1"
DESCRIBE_CLUSTER_RESULT = f'{{"cluster": "{CLUSTER_NAME}"}}'
DESCRIBE_NODEGROUP_RESULT = f'{{"nodegroup": "{NODEGROUP_NAME}"}}'
EMPTY_CLUSTER = '{"cluster": {}}'
EMPTY_NODEGROUP = '{"nodegroup": {}}'
NAME_LIST = ["foo", "bar", "baz", "qux"]
CREATE_CLUSTER_KWARGS = {"version": "1.22"}
CREATE_FARGATE_PROFILE_KWARGS = {"tags": {"hello": "world"}}
CREATE_NODEGROUP_KWARGS = {
    "capacityType": "ON_DEMAND",
    "instanceTypes": "t4g.large",
}


class ClusterParams(TypedDict):
    cluster_name: str
    cluster_role_arn: str
    resources_vpc_config: dict[Any, Any]


class NodeGroupParams(TypedDict):
    nodegroup_name: str
    nodegroup_role_arn: str


class BaseFargateProfileParams(TypedDict):
    fargate_profile_name: str
    fargate_pod_execution_role_arn: str
    fargate_selectors: list[Any]


class CreateFargateProfileParams(TypedDict):
    cluster_name: str
    pod_execution_role_arn: str
    selectors: list[Any]
    fargate_profile_name: str


class CreateNodegroupParams(TypedDict):
    cluster_name: str
    nodegroup_name: str
    nodegroup_subnets: list[str]
    nodegroup_role_arn: str


class TestEksCreateClusterOperator:
    def setup_method(self) -> None:
        # Parameters which are needed to create a cluster.
        self.create_cluster_params = ClusterParams(
            cluster_name=CLUSTER_NAME,
            cluster_role_arn=ROLE_ARN[1],
            resources_vpc_config=RESOURCES_VPC_CONFIG[1],
        )
        self.nodegroup_setup()
        self.fargate_profile_setup()

    def nodegroup_setup(self) -> None:
        # Parameters which are added to the cluster parameters
        # when creating both the cluster and nodegroup together.
        self.base_nodegroup_params = NodeGroupParams(
            nodegroup_name=NODEGROUP_NAME,
            nodegroup_role_arn=NODEROLE_ARN[1],
        )

        # Parameters expected to be passed in the CreateNodegroup hook call.
        self.create_nodegroup_params = dict(
            **self.base_nodegroup_params,
            cluster_name=CLUSTER_NAME,
            subnets=SUBNET_IDS,
        )

        self.create_cluster_operator_with_nodegroup = EksCreateClusterOperator(
            task_id=TASK_ID,
            **self.create_cluster_params,
            **self.base_nodegroup_params,
        )

    def fargate_profile_setup(self) -> None:
        # Parameters which are added to the cluster parameters
        # when creating both the cluster and Fargate profile together.
        self.base_fargate_profile_params = BaseFargateProfileParams(
            fargate_profile_name=FARGATE_PROFILE_NAME,
            fargate_pod_execution_role_arn=POD_EXECUTION_ROLE_ARN[1],
            fargate_selectors=SELECTORS[1],
        )

        # Parameters expected to be passed in the CreateFargateProfile hook call.
        self.create_fargate_profile_params = dict(
            **self.base_fargate_profile_params,
            cluster_name=CLUSTER_NAME,
        )

        self.create_cluster_operator_with_fargate_profile = EksCreateClusterOperator(
            task_id=TASK_ID,
            **self.create_cluster_params,
            **self.base_fargate_profile_params,
            compute="fargate",
        )

    @pytest.mark.parametrize(
        "create_cluster_kwargs",
        [
            pytest.param(None, id="without cluster kwargs"),
            pytest.param(CREATE_CLUSTER_KWARGS, id="with cluster kwargs"),
        ],
    )
    @mock.patch.object(EksHook, "create_cluster")
    @mock.patch.object(EksHook, "create_nodegroup")
    def test_execute_create_cluster(self, mock_create_nodegroup, mock_create_cluster, create_cluster_kwargs):
        op_kwargs = {**self.create_cluster_params, "compute": None}
        if create_cluster_kwargs:
            op_kwargs["create_cluster_kwargs"] = create_cluster_kwargs
            parameters = {**self.create_cluster_params, **create_cluster_kwargs}
        else:
            assert "create_cluster_kwargs" not in op_kwargs
            parameters = self.create_cluster_params

        operator = EksCreateClusterOperator(task_id=TASK_ID, **op_kwargs)
        operator.execute({})
        mock_create_cluster.assert_called_with(**convert_keys(parameters))
        mock_create_nodegroup.assert_not_called()

    @pytest.mark.parametrize(
        "create_cluster_kwargs",
        [
            pytest.param(None, id="without cluster kwargs"),
            pytest.param(CREATE_CLUSTER_KWARGS, id="with cluster kwargs"),
        ],
    )
    @mock.patch.object(Waiter, "wait")
    @mock.patch.object(EksHook, "create_cluster")
    @mock.patch.object(EksHook, "create_nodegroup")
    def test_execute_create_cluster_with_wait(
        self, mock_create_nodegroup, mock_create_cluster, mock_waiter, create_cluster_kwargs
    ):
        op_kwargs = {**self.create_cluster_params, "compute": None}
        if create_cluster_kwargs:
            op_kwargs["create_cluster_kwargs"] = create_cluster_kwargs
            parameters = {**self.create_cluster_params, **create_cluster_kwargs}
        else:
            assert "create_cluster_kwargs" not in op_kwargs
            parameters = self.create_cluster_params

        operator = EksCreateClusterOperator(task_id=TASK_ID, **op_kwargs, wait_for_completion=True)
        operator.execute({})
        mock_create_cluster.assert_called_with(**convert_keys(parameters))
        mock_create_nodegroup.assert_not_called()
        mock_waiter.assert_called_once_with(
            mock.ANY,
            name=CLUSTER_NAME,
            WaiterConfig={"Delay": mock.ANY, "MaxAttempts": mock.ANY},
        )
        assert_expected_waiter_type(mock_waiter, "ClusterActive")

    @mock.patch.object(Waiter, "wait")
    @mock.patch.object(EksHook, "get_cluster_state")
    @mock.patch.object(EksHook, "create_cluster")
    @mock.patch.object(EksHook, "create_nodegroup")
    def test_execute_when_called_with_nodegroup_creates_both(
        self, mock_create_nodegroup, mock_create_cluster, mock_cluster_state, mock_waiter
    ):
        mock_cluster_state.return_value = ClusterStates.ACTIVE

        self.create_cluster_operator_with_nodegroup.execute({})

        mock_create_cluster.assert_called_once_with(**convert_keys(self.create_cluster_params))
        mock_create_nodegroup.assert_called_once_with(**convert_keys(self.create_nodegroup_params))
        mock_waiter.assert_called_once_with(
            mock.ANY,
            name=CLUSTER_NAME,
            WaiterConfig={"Delay": mock.ANY, "MaxAttempts": mock.ANY},
        )
        assert_expected_waiter_type(mock_waiter, "ClusterActive")

    @mock.patch.object(Waiter, "wait")
    @mock.patch.object(EksHook, "get_cluster_state")
    @mock.patch.object(EksHook, "create_cluster")
    @mock.patch.object(EksHook, "create_nodegroup")
    def test_execute_with_wait_when_called_with_nodegroup_creates_both(
        self, mock_create_nodegroup, mock_create_cluster, mock_cluster_state, mock_waiter
    ):
        mock_cluster_state.return_value = ClusterStates.ACTIVE
        self.create_cluster_operator_with_nodegroup.wait_for_completion = True

        self.create_cluster_operator_with_nodegroup.execute({})

        mock_create_cluster.assert_called_once_with(**convert_keys(self.create_cluster_params))
        mock_create_nodegroup.assert_called_once_with(**convert_keys(self.create_nodegroup_params))
        # Calls waiter once for the cluster and once for the nodegroup.
        assert mock_waiter.call_count == 2
        mock_waiter.assert_called_with(
            mock.ANY,
            clusterName=CLUSTER_NAME,
            nodegroupName=NODEGROUP_NAME,
            WaiterConfig={"MaxAttempts": mock.ANY},
        )
        assert_expected_waiter_type(mock_waiter, "NodegroupActive")

    @mock.patch.object(Waiter, "wait")
    @mock.patch.object(EksHook, "get_cluster_state")
    @mock.patch.object(EksHook, "create_cluster")
    @mock.patch.object(EksHook, "create_fargate_profile")
    def test_execute_when_called_with_fargate_creates_both(
        self, mock_create_fargate_profile, mock_create_cluster, mock_cluster_state, mock_waiter
    ):
        mock_cluster_state.return_value = ClusterStates.ACTIVE

        self.create_cluster_operator_with_fargate_profile.execute({})

        mock_create_cluster.assert_called_once_with(**convert_keys(self.create_cluster_params))
        mock_create_fargate_profile.assert_called_once_with(
            **convert_keys(self.create_fargate_profile_params)
        )
        mock_waiter.assert_called_once_with(
            mock.ANY,
            name=CLUSTER_NAME,
            WaiterConfig={"Delay": mock.ANY, "MaxAttempts": mock.ANY},
        )
        assert_expected_waiter_type(mock_waiter, "ClusterActive")

    @mock.patch.object(Waiter, "wait")
    @mock.patch.object(EksHook, "get_cluster_state")
    @mock.patch.object(EksHook, "create_cluster")
    @mock.patch.object(EksHook, "create_fargate_profile")
    def test_execute_with_wait_when_called_with_fargate_creates_both(
        self, mock_create_fargate_profile, mock_create_cluster, mock_cluster_state, mock_waiter
    ):
        mock_cluster_state.return_value = ClusterStates.ACTIVE
        self.create_cluster_operator_with_fargate_profile.wait_for_completion = True

        self.create_cluster_operator_with_fargate_profile.execute({})

        mock_create_cluster.assert_called_once_with(**convert_keys(self.create_cluster_params))
        mock_create_fargate_profile.assert_called_once_with(
            **convert_keys(self.create_fargate_profile_params)
        )
        # Calls waiter once for the cluster and once for the nodegroup.
        assert mock_waiter.call_count == 2
        mock_waiter.assert_called_with(
            mock.ANY,
            clusterName=CLUSTER_NAME,
            fargateProfileName=FARGATE_PROFILE_NAME,
            WaiterConfig={"MaxAttempts": mock.ANY},
        )
        assert_expected_waiter_type(mock_waiter, "FargateProfileActive")

    def test_invalid_compute_value(self):
        invalid_compute = EksCreateClusterOperator(
            task_id=TASK_ID,
            **self.create_cluster_params,
            compute="infinite",
        )

        with pytest.raises(ValueError, match="Provided compute type is not supported."):
            invalid_compute.execute({})

    def test_nodegroup_compute_missing_nodegroup_role_arn(self):
        missing_nodegroup_role_arn = EksCreateClusterOperator(
            task_id=TASK_ID,
            **self.create_cluster_params,
            compute="nodegroup",
        )

        with pytest.raises(
            ValueError,
            match="Creating an Amazon EKS managed node groups requires nodegroup_role_arn to be passed in.",
        ):
            missing_nodegroup_role_arn.execute({})

    def test_fargate_compute_missing_fargate_pod_execution_role_arn(self):
        missing_fargate_pod_execution_role_arn = EksCreateClusterOperator(
            task_id=TASK_ID,
            **self.create_cluster_params,
            compute="fargate",
        )

        with pytest.raises(
            ValueError,
            match="Creating an AWS Fargate profiles requires fargate_pod_execution_role_arn to be passed in.",
        ):
            missing_fargate_pod_execution_role_arn.execute({})

    def test_init_with_region(self):
        with pytest.warns(AirflowProviderDeprecationWarning) as m:
            m.operator = EksCreateClusterOperator(
                task_id=TASK_ID,
                **self.create_cluster_params,
                compute=None,
                region="us-east-2",
            )
        assert m.operator.region_name == "us-east-2"

    @mock.patch.object(EksHook, "create_cluster")
    def test_eks_create_cluster_short_circuit_early(self, mock_create_cluster, caplog):
        mock_create_cluster.return_value = None
        eks_create_cluster_operator = EksCreateClusterOperator(
            task_id=TASK_ID,
            **self.create_cluster_params,
            compute=None,
            wait_for_completion=False,
            deferrable=False,
        )
        eks_create_cluster_operator.execute({})
        assert len(caplog.records) == 0

    @mock.patch.object(EksHook, "create_cluster")
    def test_eks_create_cluster_with_deferrable(self, mock_create_cluster, caplog):
        mock_create_cluster.return_value = None

        eks_create_cluster_operator = EksCreateClusterOperator(
            task_id=TASK_ID,
            **self.create_cluster_params,
            compute=None,
            wait_for_completion=False,
            deferrable=True,
        )
        with pytest.raises(TaskDeferred):
            eks_create_cluster_operator.execute({})
        assert "Waiting for EKS Cluster to provision. This will take some time." in caplog.messages

    def test_template_fields(self):
        op = EksCreateClusterOperator(
            task_id=TASK_ID, **self.create_cluster_params, compute="fargate", region_name="us-east-1"
        )

        validate_template_fields(op)


class TestEksCreateFargateProfileOperator:
    def setup_method(self) -> None:
        self.create_fargate_profile_params = CreateFargateProfileParams(
            cluster_name=CLUSTER_NAME,
            pod_execution_role_arn=POD_EXECUTION_ROLE_ARN[1],
            selectors=SELECTORS[1],
            fargate_profile_name=FARGATE_PROFILE_NAME,
        )

    @pytest.mark.parametrize(
        "create_fargate_profile_kwargs",
        [
            pytest.param(None, id="without fargate profile kwargs"),
            pytest.param(CREATE_FARGATE_PROFILE_KWARGS, id="with fargate profile kwargs"),
        ],
    )
    @mock.patch.object(Waiter, "wait")
    @mock.patch.object(EksHook, "create_fargate_profile")
    def test_execute_when_fargate_profile_does_not_already_exist(
        self, mock_create_fargate_profile, mock_waiter, create_fargate_profile_kwargs
    ):
        op_kwargs = {**self.create_fargate_profile_params}
        if create_fargate_profile_kwargs:
            op_kwargs["create_fargate_profile_kwargs"] = create_fargate_profile_kwargs
            parameters = {**self.create_fargate_profile_params, **create_fargate_profile_kwargs}
        else:
            assert "create_fargate_profile_kwargs" not in op_kwargs
            parameters = self.create_fargate_profile_params

        operator = EksCreateFargateProfileOperator(task_id=TASK_ID, **op_kwargs)
        operator.execute({})
        mock_create_fargate_profile.assert_called_with(**convert_keys(parameters))
        mock_waiter.assert_not_called()

    @pytest.mark.parametrize(
        "create_fargate_profile_kwargs",
        [
            pytest.param(None, id="without fargate profile kwargs"),
            pytest.param(CREATE_FARGATE_PROFILE_KWARGS, id="with fargate profile kwargs"),
        ],
    )
    @mock.patch.object(Waiter, "wait")
    @mock.patch.object(EksHook, "create_fargate_profile")
    def test_execute_with_wait_when_fargate_profile_does_not_already_exist(
        self, mock_create_fargate_profile, mock_waiter, create_fargate_profile_kwargs
    ):
        op_kwargs = {**self.create_fargate_profile_params}
        if create_fargate_profile_kwargs:
            op_kwargs["create_fargate_profile_kwargs"] = create_fargate_profile_kwargs
            parameters = {**self.create_fargate_profile_params, **create_fargate_profile_kwargs}
        else:
            assert "create_fargate_profile_kwargs" not in op_kwargs
            parameters = self.create_fargate_profile_params

        operator = EksCreateFargateProfileOperator(task_id=TASK_ID, **op_kwargs, wait_for_completion=True)
        operator.execute({})
        mock_create_fargate_profile.assert_called_with(**convert_keys(parameters))
        mock_waiter.assert_called_with(
            mock.ANY,
            clusterName=CLUSTER_NAME,
            fargateProfileName=FARGATE_PROFILE_NAME,
            WaiterConfig={"MaxAttempts": mock.ANY},
        )
        assert_expected_waiter_type(mock_waiter, "FargateProfileActive")

    @mock.patch.object(EksHook, "create_fargate_profile")
    def test_create_fargate_profile_deferrable(self, _):
        op_kwargs = {**self.create_fargate_profile_params}
        operator = EksCreateFargateProfileOperator(
            task_id=TASK_ID,
            **op_kwargs,
            deferrable=True,
        )
        with pytest.raises(TaskDeferred) as exc:
            operator.execute({})
        assert isinstance(exc.value.trigger, EksCreateFargateProfileTrigger), (
            "Trigger is not a EksCreateFargateProfileTrigger"
        )

    def test_template_fields(self):
        op = EksCreateFargateProfileOperator(task_id=TASK_ID, **self.create_fargate_profile_params)

        validate_template_fields(op)

    def test_init_with_region(self):
        with pytest.warns(AirflowProviderDeprecationWarning) as m:
            m.operator = EksCreateFargateProfileOperator(
                task_id=TASK_ID,
                **self.create_fargate_profile_params,
                region="us-east-2",
            )
        assert m.operator.region_name == "us-east-2"


class TestEksCreateNodegroupOperator:
    def setup_method(self) -> None:
        self.create_nodegroup_params = CreateNodegroupParams(
            cluster_name=CLUSTER_NAME,
            nodegroup_name=NODEGROUP_NAME,
            nodegroup_subnets=SUBNET_IDS,
            nodegroup_role_arn=NODEROLE_ARN[1],
        )

    @pytest.mark.parametrize(
        "create_nodegroup_kwargs",
        [
            pytest.param(None, id="without nodegroup kwargs"),
            pytest.param(CREATE_NODEGROUP_KWARGS, id="with nodegroup kwargs"),
        ],
    )
    @mock.patch.object(Waiter, "wait")
    @mock.patch.object(EksHook, "create_nodegroup")
    def test_execute_when_nodegroup_does_not_already_exist(
        self, mock_create_nodegroup, mock_waiter, create_nodegroup_kwargs
    ):
        op_kwargs = {**self.create_nodegroup_params}
        if create_nodegroup_kwargs:
            op_kwargs["create_nodegroup_kwargs"] = create_nodegroup_kwargs
            parameters = {**self.create_nodegroup_params, **create_nodegroup_kwargs}
        else:
            assert "create_nodegroup_params" not in op_kwargs
            parameters = self.create_nodegroup_params

        operator = EksCreateNodegroupOperator(task_id=TASK_ID, **op_kwargs)
        operator.execute({})
        mock_create_nodegroup.assert_called_with(**convert_keys(parameters))
        mock_waiter.assert_not_called()

        @pytest.mark.parametrize(
            "create_nodegroup_kwargs",
            [
                pytest.param(None, id="without nodegroup kwargs"),
                pytest.param(CREATE_NODEGROUP_KWARGS, id="with nodegroup kwargs"),
            ],
        )
        @mock.patch.object(Waiter, "wait")
        @mock.patch.object(EksHook, "create_nodegroup")
        def test_execute_with_wait_when_nodegroup_does_not_already_exist(
            self, mock_create_nodegroup, mock_waiter, create_nodegroup_kwargs
        ):
            op_kwargs = {**self.create_nodegroup_params}
            if create_nodegroup_kwargs:
                op_kwargs["create_nodegroup_kwargs"] = create_nodegroup_kwargs
                parameters = {**self.create_nodegroup_params, **create_nodegroup_kwargs}
            else:
                assert "create_nodegroup_params" not in op_kwargs
                parameters = self.create_nodegroup_params

            operator = EksCreateNodegroupOperator(task_id=TASK_ID, **op_kwargs, wait_for_completion=True)
            operator.execute({})
            mock_create_nodegroup.assert_called_with(**convert_keys(parameters))
            mock_waiter.assert_called_with(mock.ANY, clusterName=CLUSTER_NAME, nodegroupName=NODEGROUP_NAME)
            assert_expected_waiter_type(mock_waiter, "NodegroupActive")

    @mock.patch.object(EksHook, "create_nodegroup")
    def test_create_nodegroup_deferrable(self, mock_create_nodegroup):
        mock_create_nodegroup.return_value = True
        op_kwargs = {**self.create_nodegroup_params}
        operator = EksCreateNodegroupOperator(
            task_id=TASK_ID,
            **op_kwargs,
            deferrable=True,
        )
        with pytest.raises(TaskDeferred) as exc:
            operator.execute({})
        assert isinstance(exc.value.trigger, EksCreateNodegroupTrigger)

    def test_create_nodegroup_deferrable_versus_wait_for_completion(self):
        op_kwargs = {**self.create_nodegroup_params}
        operator = EksCreateNodegroupOperator(
            task_id=TASK_ID,
            **op_kwargs,
            deferrable=True,
            wait_for_completion=True,
        )
        assert operator.wait_for_completion is False
        operator = EksCreateNodegroupOperator(
            task_id=TASK_ID,
            **op_kwargs,
            deferrable=False,
            wait_for_completion=True,
        )
        assert operator.wait_for_completion is True

    def test_template_fields(self):
        op_kwargs = {**self.create_nodegroup_params}
        op = EksCreateNodegroupOperator(task_id=TASK_ID, **op_kwargs)

        validate_template_fields(op)

    def test_init_with_region(self):
        op_kwargs = {**self.create_nodegroup_params}

        with pytest.warns(AirflowProviderDeprecationWarning) as m:
            m.operator = EksCreateNodegroupOperator(
                task_id=TASK_ID,
                **op_kwargs,
                region="us-east-2",
            )
        assert m.operator.region_name == "us-east-2"


class TestEksDeleteClusterOperator:
    def setup_method(self) -> None:
        self.cluster_name: str = CLUSTER_NAME

        self.delete_cluster_operator = EksDeleteClusterOperator(
            task_id=TASK_ID, cluster_name=self.cluster_name
        )

    @mock.patch.object(Waiter, "wait")
    @mock.patch.object(EksHook, "list_nodegroups")
    @mock.patch.object(EksHook, "delete_cluster")
    def test_existing_cluster_not_in_use(self, mock_delete_cluster, mock_list_nodegroups, mock_waiter):
        mock_list_nodegroups.return_value = []
        self.delete_cluster_operator.execute({})
        mock_delete_cluster.assert_called_once_with(name=self.cluster_name)
        mock_waiter.assert_not_called()

    @mock.patch.object(Waiter, "wait")
    @mock.patch.object(EksHook, "list_nodegroups")
    @mock.patch.object(EksHook, "delete_cluster")
    def test_existing_cluster_not_in_use_with_wait(
        self, mock_delete_cluster, mock_list_nodegroups, mock_waiter
    ):
        mock_list_nodegroups.return_value = []
        self.delete_cluster_operator.wait_for_completion = True

        self.delete_cluster_operator.execute({})

        mock_list_nodegroups.assert_not_called()
        mock_delete_cluster.assert_called_once_with(name=self.cluster_name)
        mock_waiter.assert_called_with(mock.ANY, name=CLUSTER_NAME)
        assert_expected_waiter_type(mock_waiter, "ClusterDeleted")

    def test_eks_delete_cluster_operator_with_deferrable(self):
        self.delete_cluster_operator.deferrable = True
        with pytest.raises(TaskDeferred):
            self.delete_cluster_operator.execute({})

    def test_template_fields(self):
        validate_template_fields(self.delete_cluster_operator)

    def test_init_with_region(self):
        with pytest.warns(AirflowProviderDeprecationWarning) as m:
            m.operator = EksDeleteClusterOperator(
                task_id=TASK_ID, cluster_name=self.cluster_name, region="us-east-2"
            )
        assert m.operator.region_name == "us-east-2"


class TestEksDeleteNodegroupOperator:
    def setup_method(self) -> None:
        self.cluster_name: str = CLUSTER_NAME
        self.nodegroup_name: str = NODEGROUP_NAME

        self.delete_nodegroup_operator = EksDeleteNodegroupOperator(
            task_id=TASK_ID, cluster_name=self.cluster_name, nodegroup_name=self.nodegroup_name
        )

    @mock.patch.object(Waiter, "wait")
    @mock.patch.object(EksHook, "delete_nodegroup")
    def test_existing_nodegroup(self, mock_delete_nodegroup, mock_waiter):
        self.delete_nodegroup_operator.execute({})

        mock_delete_nodegroup.assert_called_once_with(
            clusterName=self.cluster_name, nodegroupName=self.nodegroup_name
        )
        mock_waiter.assert_not_called()

    @mock.patch.object(Waiter, "wait")
    @mock.patch.object(EksHook, "delete_nodegroup")
    def test_existing_nodegroup_with_wait(self, mock_delete_nodegroup, mock_waiter):
        self.delete_nodegroup_operator.wait_for_completion = True

        self.delete_nodegroup_operator.execute({})

        mock_delete_nodegroup.assert_called_once_with(
            clusterName=self.cluster_name, nodegroupName=self.nodegroup_name
        )
        mock_waiter.assert_called_with(mock.ANY, clusterName=CLUSTER_NAME, nodegroupName=NODEGROUP_NAME)
        assert_expected_waiter_type(mock_waiter, "NodegroupDeleted")

    def test_template_fields(self):
        validate_template_fields(self.delete_nodegroup_operator)

    def test_init_with_region(self):
        with pytest.warns(AirflowProviderDeprecationWarning) as m:
            m.operator = EksDeleteNodegroupOperator(
                task_id=TASK_ID,
                cluster_name=self.cluster_name,
                nodegroup_name=self.nodegroup_name,
                region="us-east-2",
            )
        assert m.operator.region_name == "us-east-2"


class TestEksDeleteFargateProfileOperator:
    def setup_method(self) -> None:
        self.cluster_name: str = CLUSTER_NAME
        self.fargate_profile_name: str = FARGATE_PROFILE_NAME

        self.delete_fargate_profile_operator = EksDeleteFargateProfileOperator(
            task_id=TASK_ID, cluster_name=self.cluster_name, fargate_profile_name=self.fargate_profile_name
        )

    @mock.patch.object(Waiter, "wait")
    @mock.patch.object(EksHook, "delete_fargate_profile")
    def test_existing_fargate_profile(self, mock_delete_fargate_profile, mock_waiter):
        self.delete_fargate_profile_operator.execute({})

        mock_delete_fargate_profile.assert_called_once_with(
            clusterName=self.cluster_name, fargateProfileName=self.fargate_profile_name
        )
        mock_waiter.assert_not_called()

    @mock.patch.object(Waiter, "wait")
    @mock.patch.object(EksHook, "delete_fargate_profile")
    def test_existing_fargate_profile_with_wait(self, mock_delete_fargate_profile, mock_waiter):
        self.delete_fargate_profile_operator.wait_for_completion = True

        self.delete_fargate_profile_operator.execute({})

        mock_delete_fargate_profile.assert_called_once_with(
            clusterName=self.cluster_name, fargateProfileName=self.fargate_profile_name
        )
        mock_waiter.assert_called_with(
            mock.ANY,
            clusterName=CLUSTER_NAME,
            fargateProfileName=FARGATE_PROFILE_NAME,
            WaiterConfig={"Delay": 30, "MaxAttempts": 60},
        )
        assert_expected_waiter_type(mock_waiter, "FargateProfileDeleted")

    @mock.patch.object(EksHook, "delete_fargate_profile")
    def test_delete_fargate_profile_deferrable(self, _):
        self.delete_fargate_profile_operator.deferrable = True

        with pytest.raises(TaskDeferred) as exc:
            self.delete_fargate_profile_operator.execute({})
        assert isinstance(exc.value.trigger, EksDeleteFargateProfileTrigger), (
            "Trigger is not a EksDeleteFargateProfileTrigger"
        )

    def test_template_fields(self):
        validate_template_fields(self.delete_fargate_profile_operator)

    def test_init_with_region(self):
        with pytest.warns(AirflowProviderDeprecationWarning) as m:
            m.operator = EksDeleteFargateProfileOperator(
                task_id=TASK_ID,
                cluster_name=self.cluster_name,
                fargate_profile_name=self.fargate_profile_name,
                region="us-east-2",
            )
        assert m.operator.region_name == "us-east-2"


class TestEksPodOperator:
    @mock.patch("airflow.providers.cncf.kubernetes.operators.pod.KubernetesPodOperator.execute")
    @mock.patch("airflow.providers.amazon.aws.hooks.eks.EksHook.generate_config_file")
    @mock.patch("airflow.providers.amazon.aws.hooks.eks.EksHook._secure_credential_context")
    @mock.patch("airflow.providers.amazon.aws.hooks.eks.EksHook.get_session")
    @mock.patch("airflow.providers.amazon.aws.hooks.eks.EksHook.__init__", return_value=None)
    def test_existing_nodegroup(
        self,
        mock_eks_hook,
        mock_get_session,
        mock_secure_credential_context,
        mock_generate_config_file,
        mock_k8s_pod_operator_execute,
    ):
        ti_context = mock.MagicMock(name="ti_context")

        # Mock the credential chain
        mock_session = mock.MagicMock()
        mock_credentials = mock.MagicMock()
        mock_frozen_credentials = mock.MagicMock()
        mock_frozen_credentials.access_key = "test_access_key"
        mock_frozen_credentials.secret_key = "test_secret_key"
        mock_frozen_credentials.token = "test_token"

        mock_get_session.return_value = mock_session
        mock_session.get_credentials.return_value = mock_credentials
        mock_credentials.get_frozen_credentials.return_value = mock_frozen_credentials

        # Mock the credential context manager
        mock_credentials_file = "/tmp/test_creds.aws_creds"
        mock_secure_credential_context.return_value.__enter__.return_value = mock_credentials_file

        # Mock the config file context manager
        mock_config_file = "/tmp/test_kubeconfig"
        mock_generate_config_file.return_value.__enter__.return_value = mock_config_file

        op = EksPodOperator(
            task_id="run_pod",
            pod_name="run_pod",
            cluster_name=CLUSTER_NAME,
            image="amazon/aws-cli:latest",
            cmds=["sh", "-c", "ls"],
            labels={"demo": "hello_world"},
            get_logs=True,
            # Delete the pod when it reaches its final state, or the execution is interrupted.
            on_finish_action="delete_pod",
        )
        op_return_value = op.execute(ti_context)

        # Verify all the expected calls were made
        mock_k8s_pod_operator_execute.assert_called_once_with(ti_context)
        mock_eks_hook.assert_called_once_with(aws_conn_id="aws_default", region_name=None)
        mock_get_session.assert_called_once()
        mock_session.get_credentials.assert_called_once()
        mock_credentials.get_frozen_credentials.assert_called_once()
        mock_secure_credential_context.assert_called_once_with(
            "test_access_key", "test_secret_key", "test_token"
        )
        mock_generate_config_file.assert_called_once_with(
            eks_cluster_name=CLUSTER_NAME, pod_namespace="default", credentials_file=mock_credentials_file
        )

        assert mock_k8s_pod_operator_execute.return_value == op_return_value
        assert op.config_file == mock_config_file

    @pytest.mark.parametrize(
        ("compatible_kpo", "kwargs", "expected_attributes"),
        [
            (
                True,
                {"on_finish_action": "delete_succeeded_pod"},
                {"on_finish_action": OnFinishAction.DELETE_SUCCEEDED_POD},
            ),
            (
                # test default
                True,
                {},
                {"on_finish_action": OnFinishAction.DELETE_POD},
            ),
            (
                # test default
                False,
                {},
                {},
            ),
        ],
    )
    def test_on_finish_action_handler(self, compatible_kpo, kwargs, expected_attributes):
        kpo_init_args_mock = mock.MagicMock(**{"parameters": ["on_finish_action"] if compatible_kpo else []})

        with mock.patch("inspect.signature", return_value=kpo_init_args_mock):
            op = EksPodOperator(
                task_id="run_pod",
                pod_name="run_pod",
                cluster_name=CLUSTER_NAME,
                image="amazon/aws-cli:latest",
                cmds=["sh", "-c", "ls"],
                labels={"demo": "hello_world"},
                get_logs=True,
                **kwargs,
            )
            for expected_attr in expected_attributes:
                assert op.__getattribute__(expected_attr) == expected_attributes[expected_attr]

    def test_template_fields(self):
        op = EksPodOperator(
            task_id="run_pod",
            pod_name="run_pod",
            cluster_name=CLUSTER_NAME,
            image="amazon/aws-cli:latest",
            cmds=["sh", "-c", "ls"],
            labels={"demo": "hello_world"},
            get_logs=True,
            on_finish_action="delete_pod",
        )

        validate_template_fields(op)

    @mock.patch("airflow.providers.cncf.kubernetes.operators.pod.KubernetesPodOperator.trigger_reentry")
    @mock.patch("airflow.providers.amazon.aws.hooks.eks.EksHook.generate_config_file")
    @mock.patch("airflow.providers.amazon.aws.hooks.eks.EksHook._secure_credential_context")
    @mock.patch("airflow.providers.amazon.aws.hooks.eks.EksHook.get_session")
    @mock.patch("airflow.providers.amazon.aws.hooks.eks.EksHook.__init__", return_value=None)
    def test_trigger_reentry(
        self,
        mock_eks_hook,
        mock_get_session,
        mock_secure_credential_context,
        mock_generate_config_file,
        mock_k8s_pod_operator_trigger_reentry,
    ):
        ti_context = mock.MagicMock(name="ti_context")
        event = {"eks_cluster_name": "eks_cluster_name", "namespace": "namespace"}

        # Mock the credential chain
        mock_session = mock.MagicMock()
        mock_credentials = mock.MagicMock()
        mock_frozen_credentials = mock.MagicMock()
        mock_frozen_credentials.access_key = "test_access_key"
        mock_frozen_credentials.secret_key = "test_secret_key"
        mock_frozen_credentials.token = "test_token"

        mock_get_session.return_value = mock_session
        mock_session.get_credentials.return_value = mock_credentials
        mock_credentials.get_frozen_credentials.return_value = mock_frozen_credentials

        # Mock the credential context manager
        mock_credentials_file = "/tmp/test_creds.aws_creds"
        mock_secure_credential_context.return_value.__enter__.return_value = mock_credentials_file

        # Mock the config file context manager
        mock_config_file = "/tmp/test_kubeconfig"
        mock_generate_config_file.return_value.__enter__.return_value = mock_config_file

        op = EksPodOperator(
            task_id="run_pod",
            pod_name="run_pod",
            cluster_name=CLUSTER_NAME,
            image="amazon/aws-cli:latest",
            cmds=["sh", "-c", "ls"],
            labels={"demo": "hello_world"},
            get_logs=True,
            # Delete the pod when it reaches its final state, or the execution is interrupted.
            on_finish_action="delete_pod",
        )
        op.trigger_reentry(ti_context, event)

        # Verify all the expected calls were made
        mock_k8s_pod_operator_trigger_reentry.assert_called_once_with(ti_context, event)
        mock_get_session.assert_called_once()
        mock_session.get_credentials.assert_called_once()
        mock_credentials.get_frozen_credentials.assert_called_once()
        mock_secure_credential_context.assert_called_once_with(
            "test_access_key", "test_secret_key", "test_token"
        )
        mock_generate_config_file.assert_called_once_with(
            eks_cluster_name="eks_cluster_name",
            pod_namespace="namespace",
            credentials_file=mock_credentials_file,
        )
        assert op.config_file == mock_config_file
