
import copy

import pytest
from unittest.mock import Mock, patch

from ray.autoscaler._private.aws.config import _configure_subnet, \
    _get_vpc_id_or_die, bootstrap_aws, log_to_cli, \
    DEFAULT_AMI
from ray.autoscaler._private.aws.node_provider import AWSNodeProvider
from ray.autoscaler._private.providers import _get_node_provider
import ray.tests.aws.utils.stubs as stubs
import ray.tests.aws.utils.helpers as helpers
from ray.tests.aws.utils.constants import AUX_SUBNET, DEFAULT_SUBNET, \
    DEFAULT_SG_AUX_SUBNET, DEFAULT_SG, DEFAULT_SG_DUAL_GROUP_RULES, \
    DEFAULT_SG_WITH_RULES_AUX_SUBNET, AUX_SG, \
    DEFAULT_SG_WITH_RULES, DEFAULT_SG_WITH_NAME, \
    DEFAULT_SG_WITH_NAME_AND_RULES, CUSTOM_IN_BOUND_RULES, \
    DEFAULT_KEY_PAIR, DEFAULT_INSTANCE_PROFILE, DEFAULT_CLUSTER_NAME, \
    DEFAULT_LT


def test_use_subnets_in_only_one_vpc(iam_client_stub, ec2_client_stub):
    """
    This test validates that when bootstrap_aws populates the SubnetIds field,
    all of the subnets used belong to the same VPC, and that a SecurityGroup
    in that VPC is correctly configured.
    Also validates that head IAM role is correctly filled.
    """
    stubs.configure_iam_role_default(iam_client_stub)
    stubs.configure_key_pair_default(ec2_client_stub)

    # Add a response with a thousand subnets all in different VPCs.
    # After filtering, only subnet in one particular VPC should remain.
    # Thus SubnetIds for each available node type should end up as
    # being length-one lists after the bootstrap_config.
    stubs.describe_a_thousand_subnets_in_different_vpcs(ec2_client_stub)

    # describe the subnet in use while determining its vpc
    stubs.describe_subnets_echo(ec2_client_stub, DEFAULT_SUBNET)
    # given no existing security groups within the VPC...
    stubs.describe_no_security_groups(ec2_client_stub)
    # expect to create a security group on the VPC
    stubs.create_sg_echo(ec2_client_stub, DEFAULT_SG)
    # expect new security group details to be retrieved after creation
    stubs.describe_sgs_on_vpc(
        ec2_client_stub,
        [DEFAULT_SUBNET["VpcId"]],
        [DEFAULT_SG],
    )

    # given no existing default security group inbound rules...
    # expect to authorize all default inbound rules
    stubs.authorize_sg_ingress(
        ec2_client_stub,
        DEFAULT_SG_WITH_RULES,
    )

    # expect another call to describe the above security group while checking
    # a second time if it has ip_permissions set ("if not sg.ip_permissions")
    stubs.describe_an_sg_2(
        ec2_client_stub,
        DEFAULT_SG_WITH_RULES,
    )

    # given our mocks and an example config file as input...
    # expect the config to be loaded, validated, and bootstrapped successfully
    config = helpers.bootstrap_aws_example_config_file("example-full.yaml")
    _get_vpc_id_or_die.cache_clear()

    # We've filtered down to only one subnet id -- only one of the thousand
    # subnets generated by ec2.subnets.all() belongs to the right VPC.
    for node_type in config["available_node_types"].values():
        node_config = node_type["node_config"]
        assert node_config["SubnetIds"] == [DEFAULT_SUBNET["SubnetId"]]
        assert node_config["SecurityGroupIds"] == [DEFAULT_SG["GroupId"]]


def test_create_sg_different_vpc_same_rules(iam_client_stub, ec2_client_stub):
    # use default stubs to skip ahead to security group configuration
    stubs.skip_to_configure_sg(ec2_client_stub, iam_client_stub)

    # given head and worker nodes with custom subnets defined...
    # expect to first describe the worker subnet ID
    stubs.describe_subnets_echo(ec2_client_stub, AUX_SUBNET)
    # expect to second describe the head subnet ID
    stubs.describe_subnets_echo(ec2_client_stub, DEFAULT_SUBNET)
    # given no existing security groups within the VPC...
    stubs.describe_no_security_groups(ec2_client_stub)
    # expect to first create a security group on the worker node VPC
    stubs.create_sg_echo(ec2_client_stub, DEFAULT_SG_AUX_SUBNET)
    # expect new worker security group details to be retrieved after creation
    stubs.describe_sgs_on_vpc(
        ec2_client_stub,
        [AUX_SUBNET["VpcId"]],
        [DEFAULT_SG_AUX_SUBNET],
    )
    # expect to second create a security group on the head node VPC
    stubs.create_sg_echo(ec2_client_stub, DEFAULT_SG)
    # expect new head security group details to be retrieved after creation
    stubs.describe_sgs_on_vpc(
        ec2_client_stub,
        [DEFAULT_SUBNET["VpcId"]],
        [DEFAULT_SG],
    )

    # given no existing default head security group inbound rules...
    # expect to authorize all default head inbound rules
    stubs.authorize_sg_ingress(
        ec2_client_stub,
        DEFAULT_SG_DUAL_GROUP_RULES,
    )
    # given no existing default worker security group inbound rules...
    # expect to authorize all default worker inbound rules
    stubs.authorize_sg_ingress(
        ec2_client_stub,
        DEFAULT_SG_WITH_RULES_AUX_SUBNET,
    )

    # given our mocks and an example config file as input...
    # expect the config to be loaded, validated, and bootstrapped successfully
    config = helpers.bootstrap_aws_example_config_file("example-subnets.yaml")

    # expect the bootstrapped config to show different head and worker security
    # groups residing on different subnets
    for node_type_key, node_type in config["available_node_types"].items():
        node_config = node_type["node_config"]
        security_group_ids = node_config["SecurityGroupIds"]
        subnet_ids = node_config["SubnetIds"]
        if node_type_key == config["head_node_type"]:
            assert security_group_ids == [DEFAULT_SG["GroupId"]]
            assert subnet_ids == [DEFAULT_SUBNET["SubnetId"]]
        else:
            assert security_group_ids == [AUX_SG["GroupId"]]
            assert subnet_ids == [AUX_SUBNET["SubnetId"]]

    # expect no pending responses left in IAM or EC2 client stub queues
    iam_client_stub.assert_no_pending_responses()
    ec2_client_stub.assert_no_pending_responses()


def test_create_sg_with_custom_inbound_rules_and_name(iam_client_stub,
                                                      ec2_client_stub):
    # use default stubs to skip ahead to security group configuration
    stubs.skip_to_configure_sg(ec2_client_stub, iam_client_stub)

    # expect to describe the head subnet ID
    stubs.describe_subnets_echo(ec2_client_stub, DEFAULT_SUBNET)
    # given no existing security groups within the VPC...
    stubs.describe_no_security_groups(ec2_client_stub)
    # expect to create a security group on the head node VPC
    stubs.create_sg_echo(ec2_client_stub, DEFAULT_SG_WITH_NAME)
    # expect new head security group details to be retrieved after creation
    stubs.describe_sgs_on_vpc(
        ec2_client_stub,
        [DEFAULT_SUBNET["VpcId"]],
        [DEFAULT_SG_WITH_NAME],
    )

    # given custom existing default head security group inbound rules...
    # expect to authorize both default and custom inbound rules
    stubs.authorize_sg_ingress(
        ec2_client_stub,
        DEFAULT_SG_WITH_NAME_AND_RULES,
    )

    # given the prior modification to the head security group...
    # expect the next read of a head security group property to reload it
    stubs.describe_sg_echo(ec2_client_stub, DEFAULT_SG_WITH_NAME_AND_RULES)

    _get_vpc_id_or_die.cache_clear()
    # given our mocks and an example config file as input...
    # expect the config to be loaded, validated, and bootstrapped successfully
    config = helpers.bootstrap_aws_example_config_file(
        "example-security-group.yaml")

    # expect the bootstrapped config to have the custom security group...
    # name and in bound rules
    assert config["provider"]["security_group"][
               "GroupName"] == DEFAULT_SG_WITH_NAME_AND_RULES["GroupName"]
    assert config["provider"]["security_group"][
               "IpPermissions"] == CUSTOM_IN_BOUND_RULES

    # expect no pending responses left in IAM or EC2 client stub queues
    iam_client_stub.assert_no_pending_responses()
    ec2_client_stub.assert_no_pending_responses()


def test_subnet_given_head_and_worker_sg(iam_client_stub, ec2_client_stub):
    stubs.configure_iam_role_default(iam_client_stub)
    stubs.configure_key_pair_default(ec2_client_stub)

    # list a security group and a thousand subnets in different vpcs
    stubs.describe_a_security_group(ec2_client_stub, DEFAULT_SG)
    stubs.describe_a_thousand_subnets_in_different_vpcs(ec2_client_stub)

    config = helpers.bootstrap_aws_example_config_file(
        "example-head-and-worker-security-group.yaml")

    # check that just the single subnet in the right vpc is filled
    for node_type in config["available_node_types"].values():
        node_config = node_type["node_config"]
        assert node_config["SubnetIds"] == [DEFAULT_SUBNET["SubnetId"]]

    # expect no pending responses left in IAM or EC2 client stub queues
    iam_client_stub.assert_no_pending_responses()
    ec2_client_stub.assert_no_pending_responses()


def test_fills_out_amis_and_iam(iam_client_stub, ec2_client_stub):
    # Setup stubs to mock out boto3
    stubs.configure_iam_role_default(iam_client_stub)
    stubs.configure_key_pair_default(ec2_client_stub)
    stubs.describe_a_security_group(ec2_client_stub, DEFAULT_SG)
    stubs.configure_subnet_default(ec2_client_stub)

    config = helpers.load_aws_example_config_file("example-full.yaml")
    head_node_config = config["available_node_types"]["ray.head.default"][
        "node_config"]
    worker_node_config = config["available_node_types"]["ray.worker.default"][
        "node_config"]

    del head_node_config["ImageId"]
    del worker_node_config["ImageId"]

    # Pass in SG for stub to work
    head_node_config["SecurityGroupIds"] = ["sg-1234abcd"]
    worker_node_config["SecurityGroupIds"] = ["sg-1234abcd"]

    defaults_filled = bootstrap_aws(config)

    ami = DEFAULT_AMI.get(defaults_filled.get("provider", {}).get("region"))

    for node_type in defaults_filled["available_node_types"].values():
        node_config = node_type["node_config"]
        assert node_config.get("ImageId") == ami

    # Correctly configured IAM role
    assert (defaults_filled["head_node"]["IamInstanceProfile"] == {
        "Arn": DEFAULT_INSTANCE_PROFILE["Arn"]
    })
    # Workers of the head's type do not get the IAM role.
    head_type = config["head_node_type"]
    assert "IamInstanceProfile" not in defaults_filled["available_node_types"][
        head_type]

    iam_client_stub.assert_no_pending_responses()
    ec2_client_stub.assert_no_pending_responses()


def test_iam_already_configured(iam_client_stub, ec2_client_stub):
    """
    Checks that things work as expected when IAM role is supplied by user.
    """
    stubs.configure_key_pair_default(ec2_client_stub)
    stubs.describe_a_security_group(ec2_client_stub, DEFAULT_SG)
    stubs.configure_subnet_default(ec2_client_stub)

    config = helpers.load_aws_example_config_file("example-full.yaml")
    head_node_config = config["available_node_types"]["ray.head.default"][
        "node_config"]
    worker_node_config = config["available_node_types"]["ray.worker.default"][
        "node_config"]

    head_node_config["IamInstanceProfile"] = "mock_profile"

    # Pass in SG for stub to work
    head_node_config["SecurityGroupIds"] = ["sg-1234abcd"]
    worker_node_config["SecurityGroupIds"] = ["sg-1234abcd"]

    defaults_filled = bootstrap_aws(config)
    filled_head = defaults_filled["available_node_types"]["ray.head.default"][
        "node_config"]
    assert filled_head["IamInstanceProfile"] == "mock_profile"
    assert "IamInstanceProfile" not in defaults_filled["head_node"]

    iam_client_stub.assert_no_pending_responses()
    ec2_client_stub.assert_no_pending_responses()


def test_create_sg_multinode(iam_client_stub, ec2_client_stub):
    """
    Test AWS Bootstrap logic when config being bootstrapped has the
    following properties:
    (1) auth config does not specify ssh key path
    (2) available_node_types is provided
    (3) security group name and ip permissions set in provider field
    (4) Available node types have SubnetIds field set and this
        field is of form SubnetIds: [subnet-xxxxx].
        Both node types specify the same subnet-xxxxx.
    Tests creation of a security group and key pair under these conditions.
    """

    # Generate a config of the desired form.
    subnet_id = DEFAULT_SUBNET["SubnetId"]

    # security group info to go in provider field
    provider_data = helpers.load_aws_example_config_file(
        "example-security-group.yaml")["provider"]

    # a multi-node-type config -- will add head/worker stuff and security group
    # info to this.
    base_config = helpers.load_aws_example_config_file("example-full.yaml")

    config = copy.deepcopy(base_config)
    # Add security group data
    config["provider"] = provider_data
    # Add head and worker fields.
    head_node_config = config["available_node_types"]["ray.head.default"][
        "node_config"]
    worker_node_config = config["available_node_types"]["ray.worker.default"][
        "node_config"]
    head_node_config["SubnetIds"] = [subnet_id]
    worker_node_config["SubnetIds"] = [subnet_id]

    # Generate stubs
    stubs.configure_iam_role_default(iam_client_stub)
    stubs.configure_key_pair_default(ec2_client_stub)

    # Only one of these (the one specified in the available_node_types)
    # is in the correct vpc.
    # This list of subnets is generated by the ec2.subnets.all() call
    # and then ignored, since available_node_types already specify
    # subnet_ids.
    stubs.describe_a_thousand_subnets_in_different_vpcs(ec2_client_stub)

    # The rest of the stubbing logic is copied from
    # test_create_sg_with_custom_inbound_rules_and_name.

    # expect to describe the head subnet ID
    stubs.describe_subnets_echo(ec2_client_stub, DEFAULT_SUBNET)
    # given no existing security groups within the VPC...
    stubs.describe_no_security_groups(ec2_client_stub)
    # expect to create a security group on the head node VPC
    stubs.create_sg_echo(ec2_client_stub, DEFAULT_SG_WITH_NAME)
    # expect new head security group details to be retrieved after creation
    stubs.describe_sgs_on_vpc(
        ec2_client_stub,
        [DEFAULT_SUBNET["VpcId"]],
        [DEFAULT_SG_WITH_NAME],
    )

    # given custom existing default head security group inbound rules...
    # expect to authorize both default and custom inbound rules
    stubs.authorize_sg_ingress(
        ec2_client_stub,
        DEFAULT_SG_WITH_NAME_AND_RULES,
    )

    # given the prior modification to the head security group...
    # expect the next read of a head security group property to reload it
    stubs.describe_sg_echo(ec2_client_stub, DEFAULT_SG_WITH_NAME_AND_RULES)

    _get_vpc_id_or_die.cache_clear()

    # given our mocks and the config as input...
    # expect the config to be validated and bootstrapped successfully
    bootstrapped_config = helpers.bootstrap_aws_config(config)

    # expect the bootstrapped config to have the custom security group...
    # name and in bound rules
    assert bootstrapped_config["provider"]["security_group"][
               "GroupName"] == DEFAULT_SG_WITH_NAME_AND_RULES["GroupName"]
    assert bootstrapped_config["provider"]["security_group"][
               "IpPermissions"] == CUSTOM_IN_BOUND_RULES

    # Confirming correct security group got filled for head and workers
    sg_id = DEFAULT_SG["GroupId"]
    for node_type in bootstrapped_config["available_node_types"].values():
        node_config = node_type["node_config"]
        assert node_config["SecurityGroupIds"] == [sg_id]

    # Confirming boostrap config updates available node types with
    # default KeyName
    for node_type in bootstrapped_config["available_node_types"].values():
        node_config = node_type["node_config"]
        assert node_config["KeyName"] == DEFAULT_KEY_PAIR["KeyName"]

    # Confirm security group is in the right VPC.
    # (Doesn't really confirm anything except for the structure of this test
    # data.)
    bootstrapped_head_type = bootstrapped_config["head_node_type"]
    bootstrapped_types = bootstrapped_config["available_node_types"]
    bootstrapped_head_config = bootstrapped_types[bootstrapped_head_type][
        "node_config"]
    assert DEFAULT_SG["VpcId"] == DEFAULT_SUBNET["VpcId"]
    assert DEFAULT_SUBNET["SubnetId"] == bootstrapped_head_config["SubnetIds"][
        0]

    # ssh private key filled in
    assert "ssh_private_key" in bootstrapped_config["auth"]

    # expect no pending responses left in IAM or EC2 client stub queues
    iam_client_stub.assert_no_pending_responses()
    ec2_client_stub.assert_no_pending_responses()


def test_missing_keyname(iam_client_stub, ec2_client_stub):
    config = helpers.load_aws_example_config_file("example-full.yaml")
    config["auth"]["ssh_private_key"] = "/path/to/private/key"
    head_node_config = config["available_node_types"]["ray.head.default"][
        "node_config"]
    worker_node_config = config["available_node_types"]["ray.worker.default"][
        "node_config"]

    # Setup stubs to mock out boto3. Should fail on assertion after
    # checking KeyName/UserData.
    stubs.configure_iam_role_default(iam_client_stub)

    missing_user_data_config = copy.deepcopy(config)
    with pytest.raises(AssertionError):
        # Config specified ssh_private_key, but missing KeyName/UserData in
        # node configs
        bootstrap_aws(missing_user_data_config)

    # Pass in SG for stub to work
    head_node_config["SecurityGroupIds"] = ["sg-1234abcd"]
    worker_node_config["SecurityGroupIds"] = ["sg-1234abcd"]

    # Set UserData for both node configs
    head_node_config["UserData"] = {"someKey": "someValue"}
    worker_node_config["UserData"] = {"someKey": "someValue"}

    # Stubs to mock out boto3. Should no longer fail on assertion
    # and go on to describe security groups + configure subnet
    stubs.configure_iam_role_default(iam_client_stub)
    stubs.describe_a_security_group(ec2_client_stub, DEFAULT_SG)
    stubs.configure_subnet_default(ec2_client_stub)

    # Should work without error now that UserData is set
    bootstrap_aws(config)

    iam_client_stub.assert_no_pending_responses()
    ec2_client_stub.assert_no_pending_responses()


def test_log_to_cli(iam_client_stub, ec2_client_stub):
    config = helpers.load_aws_example_config_file("example-full.yaml")

    head_node_config = config["available_node_types"]["ray.head.default"][
        "node_config"]
    worker_node_config = config["available_node_types"]["ray.worker.default"][
        "node_config"]

    # Pass in SG for stub to work
    head_node_config["SecurityGroupIds"] = ["sg-1234abcd"]
    worker_node_config["SecurityGroupIds"] = ["sg-1234abcd"]

    stubs.configure_iam_role_default(iam_client_stub)
    stubs.configure_key_pair_default(ec2_client_stub)
    stubs.describe_a_security_group(ec2_client_stub, DEFAULT_SG)
    stubs.configure_subnet_default(ec2_client_stub)

    config = helpers.bootstrap_aws_config(config)

    # Only side-effect is to print logs to cli, called just to
    # check that it runs without error
    log_to_cli(config)
    iam_client_stub.assert_no_pending_responses()
    ec2_client_stub.assert_no_pending_responses()


def test_network_interfaces(ec2_client_stub, iam_client_stub,
                            ec2_client_stub_fail_fast,
                            ec2_client_stub_max_retries):

    # use default stubs to skip ahead to subnet configuration
    stubs.configure_iam_role_default(iam_client_stub)
    stubs.configure_key_pair_default(ec2_client_stub)

    # given the security groups associated with our network interfaces...
    sgids = ["sg-00000000", "sg-11111111", "sg-22222222", "sg-33333333"]
    security_groups = []
    suffix = 0
    for sgid in sgids:
        sg = copy.deepcopy(DEFAULT_SG)
        sg["GroupName"] += f"-{suffix}"
        sg["GroupId"] = sgid
        security_groups.append(sg)
        suffix += 1
    # expect to describe all security groups to ensure they share the same VPC
    stubs.describe_sgs_by_id(ec2_client_stub, sgids, security_groups)

    # use a default stub to skip subnet configuration
    stubs.configure_subnet_default(ec2_client_stub)

    # given our mocks and an example config file as input...
    # expect the config to be loaded, validated, and bootstrapped successfully
    config = helpers.bootstrap_aws_example_config_file(
        "example-network-interfaces.yaml")

    # instantiate a new node provider
    new_provider = _get_node_provider(
        config["provider"],
        DEFAULT_CLUSTER_NAME,
        False,
    )

    for name, node_type in config["available_node_types"].items():
        node_cfg = node_type["node_config"]
        tags = helpers.node_provider_tags(config, name)
        # given our bootstrapped node config as input to create a new node...
        # expect to first describe all stopped instances that could be reused
        stubs.describe_instances_with_any_filter_consumer(
            ec2_client_stub_max_retries)
        # given no stopped EC2 instances to reuse...
        # expect to create new nodes with the given network interface config
        stubs.run_instances_with_network_interfaces_consumer(
            ec2_client_stub_fail_fast,
            node_cfg["NetworkInterfaces"],
        )
        new_provider.create_node(node_cfg, tags, 1)

    iam_client_stub.assert_no_pending_responses()
    ec2_client_stub.assert_no_pending_responses()
    ec2_client_stub_fail_fast.assert_no_pending_responses()
    ec2_client_stub_max_retries.assert_no_pending_responses()


def test_network_interface_conflict_keys():
    # If NetworkInterfaces are defined, SubnetId and SecurityGroupIds
    # can't be specified in the same node type config.
    conflict_kv_pairs = [("SubnetId", "subnet-0000000"),
                         ("SubnetIds", ["subnet-0000000", "subnet-1111111"]),
                         ("SecurityGroupIds", ["sg-1234abcd", "sg-dcba4321"])]
    expected_error_msg = "If NetworkInterfaces are defined, subnets and " \
                         "security groups must ONLY be given in each " \
                         "NetworkInterface."
    for conflict_kv_pair in conflict_kv_pairs:
        config = helpers.load_aws_example_config_file(
            "example-network-interfaces.yaml")
        head_name = config["head_node_type"]
        head_node_cfg = config["available_node_types"][head_name][
            "node_config"]
        head_node_cfg[conflict_kv_pair[0]] = conflict_kv_pair[1]
        with pytest.raises(ValueError, match=expected_error_msg):
            helpers.bootstrap_aws_config(config)


def test_network_interface_missing_subnet():
    # If NetworkInterfaces are defined, each must have a subnet ID
    expected_error_msg = "NetworkInterfaces are defined but at least one is " \
                         "missing a subnet. Please ensure all interfaces " \
                         "have a subnet assigned."
    config = helpers.load_aws_example_config_file(
        "example-network-interfaces.yaml")
    for name, node_type in config["available_node_types"].items():
        node_cfg = node_type["node_config"]
        for network_interface_cfg in node_cfg["NetworkInterfaces"]:
            network_interface_cfg.pop("SubnetId")
            with pytest.raises(ValueError, match=expected_error_msg):
                helpers.bootstrap_aws_config(config)


def test_network_interface_missing_security_group():
    # If NetworkInterfaces are defined, each must have security groups
    expected_error_msg = "NetworkInterfaces are defined but at least one is " \
                         "missing a security group. Please ensure all " \
                         "interfaces have a security group assigned."
    config = helpers.load_aws_example_config_file(
        "example-network-interfaces.yaml")
    for name, node_type in config["available_node_types"].items():
        node_cfg = node_type["node_config"]
        for network_interface_cfg in node_cfg["NetworkInterfaces"]:
            network_interface_cfg.pop("Groups")
            with pytest.raises(ValueError, match=expected_error_msg):
                helpers.bootstrap_aws_config(config)


def test_launch_templates(ec2_client_stub, ec2_client_stub_fail_fast,
                          ec2_client_stub_max_retries):

    # given the launch template associated with our default head node type...
    # expect to first describe the default launch template by ID
    stubs.describe_launch_template_versions_by_id_default(
        ec2_client_stub, ["$Latest"])
    # given the launch template associated with our default worker node type...
    # expect to next describe the same default launch template by name
    stubs.describe_launch_template_versions_by_name_default(
        ec2_client_stub, ["2"])
    # use default stubs to skip ahead to subnet configuration
    stubs.configure_key_pair_default(ec2_client_stub)

    # given the security groups associated with our launch template...
    sgids = [DEFAULT_SG["GroupId"]]
    security_groups = [DEFAULT_SG]
    # expect to describe all security groups to ensure they share the same VPC
    stubs.describe_sgs_by_id(ec2_client_stub, sgids, security_groups)

    # use a default stub to skip subnet configuration
    stubs.configure_subnet_default(ec2_client_stub)

    # given our mocks and an example config file as input...
    # expect the config to be loaded, validated, and bootstrapped successfully
    config = helpers.bootstrap_aws_example_config_file(
        "example-launch-templates.yaml")

    # instantiate a new node provider
    new_provider = _get_node_provider(
        config["provider"],
        DEFAULT_CLUSTER_NAME,
        False,
    )

    max_count = 1
    for name, node_type in config["available_node_types"].items():
        # given our bootstrapped node config as input to create a new node...
        # expect to first describe all stopped instances that could be reused
        stubs.describe_instances_with_any_filter_consumer(
            ec2_client_stub_max_retries)
        # given no stopped EC2 instances to reuse...
        # expect to create new nodes with the given launch template config
        node_cfg = node_type["node_config"]
        stubs.run_instances_with_launch_template_consumer(
            ec2_client_stub_fail_fast, config, node_cfg, name,
            DEFAULT_LT["LaunchTemplateData"], max_count)
        tags = helpers.node_provider_tags(config, name)
        new_provider.create_node(node_cfg, tags, max_count)

    ec2_client_stub.assert_no_pending_responses()
    ec2_client_stub_fail_fast.assert_no_pending_responses()
    ec2_client_stub_max_retries.assert_no_pending_responses()


@pytest.mark.parametrize("num_on_demand_nodes", [0, 1001, 9999])
@pytest.mark.parametrize("num_spot_nodes", [0, 1001, 9999])
@pytest.mark.parametrize("stop", [True, False])
def test_terminate_nodes(num_on_demand_nodes, num_spot_nodes, stop):
    # This node makes sure that we stop or terminate all the nodes we're
    # supposed to stop or terminate when we call "terminate_nodes". This test
    # alse makes sure that we don't try to stop or terminate too many nodes in
    # a single EC2 request. By default, only 1000 nodes can be
    # stopped/terminated in one request. To terminate more nodes, we must break
    # them up into multiple smaller requests.
    #
    # "num_on_demand_nodes" is the number of on-demand nodes to stop or
    #   terminate.
    # "num_spot_nodes" is the number of on-demand nodes to terminate.
    # "stop" is True if we want to stop nodes, and False to terminate nodes.
    #   Note that spot instances are always terminated, even if "stop" is True.

    # Generate a list of unique instance ids to terminate
    on_demand_nodes = {
        "i-{:017d}".format(i)
        for i in range(num_on_demand_nodes)
    }
    spot_nodes = {
        "i-{:017d}".format(i + num_on_demand_nodes)
        for i in range(num_spot_nodes)
    }
    node_ids = list(on_demand_nodes.union(spot_nodes))

    with patch("ray.autoscaler._private.aws.node_provider.make_ec2_client"):
        provider = AWSNodeProvider(
            provider_config={
                "region": "nowhere",
                "cache_stopped_nodes": stop
            },
            cluster_name="default")

    # "_get_cached_node" is used by the AWSNodeProvider to determine whether a
    # node is a spot instance or an on-demand instance.
    def mock_get_cached_node(node_id):
        result = Mock()
        result.spot_instance_request_id = "sir-08b93456" if \
            node_id in spot_nodes else ""
        return result

    provider._get_cached_node = mock_get_cached_node

    provider.terminate_nodes(node_ids)

    stop_calls = provider.ec2.meta.client.stop_instances.call_args_list
    terminate_calls = provider.ec2.meta.client.terminate_instances \
        .call_args_list

    nodes_to_stop = set()
    nodes_to_terminate = spot_nodes

    if stop:
        nodes_to_stop.update(on_demand_nodes)
    else:
        nodes_to_terminate.update(on_demand_nodes)

    for calls, nodes_to_include_in_call in (stop_calls, nodes_to_stop), (
        terminate_calls, nodes_to_terminate):
        nodes_included_in_call = set()
        for call in calls:
            assert len(call[1]["InstanceIds"]) <= provider.max_terminate_nodes
            nodes_included_in_call.update(call[1]["InstanceIds"])

        assert nodes_to_include_in_call == nodes_included_in_call


def test_use_subnets_ordered_by_az(ec2_client_stub):
    """
    This test validates that when bootstrap_aws populates the SubnetIds field,
    the subnets are ordered the same way as availability zones.
    """
    # Add a response with a twenty subnets round-robined across the 4 AZs in
    # `us-west-2` (a,b,c,d). At the end we should only have 15 subnets, ordered
    # first from `us-west-2c`, then `us-west-2d`, then `us-west-2a`.
    stubs.describe_twenty_subnets_in_different_azs(ec2_client_stub)

    base_config = helpers.load_aws_example_config_file("example-full.yaml")
    base_config["provider"][
        "availability_zone"] = "us-west-2c,us-west-2d,us-west-2a"
    config = _configure_subnet(base_config)

    # We've filtered down to only subnets in 2c, 2d & 2a
    for node_type in config["available_node_types"].values():
        node_config = node_type["node_config"]
        assert len(node_config["SubnetIds"]) == 15
        offsets = [int(s.split("-")[1]) % 4 for s in node_config["SubnetIds"]]
        assert set(offsets[:5]) == {2}, "First 5 should be in us-west-2c"
        assert set(offsets[5:10]) == {3}, "Next 5 should be in us-west-2d"
        assert set(offsets[10:15]) == {0}, "Last 5 should be in us-west-2a"


def test_cloudwatch_dashboard_creation(cloudwatch_client_stub,
                                       ssm_client_stub):
    # create test cluster node IDs and an associated cloudwatch helper
    node_id = "i-abc"
    cloudwatch_helper = helpers.get_cloudwatch_helper(node_id)

    # given a directive to create a cluster CloudWatch Dashboard...
    # expect to make a call to create a dashboard for each node in the cluster
    stubs.put_cluster_dashboard_success(
        cloudwatch_client_stub,
        cloudwatch_helper,
    )

    # given our mocks and the example CloudWatch Dashboard config as input...
    # expect a cluster CloudWatch Dashboard to be created successfully
    cloudwatch_helper._put_cloudwatch_dashboard()
    # expect no pending responses left in the CloudWatch client stub queue
    cloudwatch_client_stub.assert_no_pending_responses()


def test_cloudwatch_alarm_creation(cloudwatch_client_stub, ssm_client_stub):
    # create test cluster node IDs and an associated cloudwatch helper
    node_id = "i-abc"
    cloudwatch_helper = helpers.get_cloudwatch_helper(node_id)

    # given a directive to update a cluster CloudWatch Alarm Config without any
    # change...
    # expect the stored the CloudWatch Alarm Config is same as local config
    cw_ssm_param_name = helpers.get_ssm_param_name(
        cloudwatch_helper.cluster_name, "alarm")
    stubs.get_param_ssm_same(ssm_client_stub, cw_ssm_param_name,
                             cloudwatch_helper, "alarm")

    # given a directive to create cluster CloudWatch alarms...
    # expect to make a call to create alarms for each node in the cluster
    stubs.put_cluster_alarms_success(cloudwatch_client_stub, cloudwatch_helper)

    # given our mocks and the example CloudWatch Alarm config as input...
    # expect cluster alarms to be created successfully
    cloudwatch_helper._put_cloudwatch_alarm()

    # expect no pending responses left in the CloudWatch client stub queue
    cloudwatch_client_stub.assert_no_pending_responses()


def test_cloudwatch_agent_update_without_change_head_node(
        ssm_client_stub, ec2_client_stub):
    # create test cluster head node ID and an associated cloudwatch helper
    node_id = "i-abc"
    is_head_node = True
    cloudwatch_helper = helpers.get_cloudwatch_helper(node_id)

    # given a directive to check for the Unified CloudWatch Agent status...
    # expect CloudWatch Agent is installed
    stubs.get_ec2_cwa_installed_tag_true(ec2_client_stub, node_id)

    # given a directive to update a cluster CloudWatch Agent Config without any
    # change...
    # expect the stored the CloudWatch Agent Config is same as local config
    cw_ssm_param_name = helpers.get_ssm_param_name(
        cloudwatch_helper.cluster_name, "agent")
    stubs.get_param_ssm_same(ssm_client_stub, cw_ssm_param_name,
                             cloudwatch_helper, "agent")

    # given our mocks and the same cloudwatch agent config as input...
    # expect no update performed on CloudWatch Agent Config
    cloudwatch_helper._update_cloudwatch_config("agent", is_head_node)


def test_cloudwatch_agent_update_with_change_head_node(ec2_client_stub,
                                                       ssm_client_stub):
    # create test cluster head node ID and an associated cloudwatch helper
    node_id = "i-abc"
    is_head_node = True
    cloudwatch_helper = helpers.get_cloudwatch_helper(node_id)

    # given a directive to check for the Unified CloudWatch Agent status...
    # expect CloudWatch Agent is installed
    stubs.get_ec2_cwa_installed_tag_true(ec2_client_stub, node_id)
    # given a directive to update a cluster CloudWatch Agent Config with new
    # changes...
    # expect the stored the CloudWatch Agent Config is different from local
    # config
    cw_ssm_param_name = helpers.get_ssm_param_name(
        cloudwatch_helper.cluster_name, "agent")
    stubs.get_param_ssm_different(ssm_client_stub, cw_ssm_param_name)

    # given an updated CloudWatch Agent Config file...
    # expect to store the new CloudWatch Agent config as an SSM parameter
    cmd_id = stubs.put_parameter_cloudwatch_config(
        ssm_client_stub, cloudwatch_helper.cluster_name, "agent")

    # given an updated CloudWatch Agent Config file...
    # expect to update the node tag equal to updated config file sha1 hash
    # to reflect the changes in config file
    stubs.update_hash_tag_success(ec2_client_stub, node_id, "agent",
                                  cloudwatch_helper)
    # given that updated CloudWatch Agent Config is put to Parameter Store...
    # expect to send an SSM command to restart CloudWatch Agent on all nodes
    cmd_id = stubs.send_command_stop_cwa(ssm_client_stub, node_id)
    # given a SSM command to stop CloudWatch Agent sent to all nodes...
    # expect to wait for the command to complete successfully on every node
    stubs.list_command_invocations_success(ssm_client_stub, node_id, cmd_id)
    cmd_id = stubs.send_command_start_cwa(ssm_client_stub, node_id,
                                          cw_ssm_param_name)
    # given a SSM command to start CloudWatch Agent sent to all nodes...
    # expect to wait for the command to complete successfully on every node
    stubs.list_command_invocations_success(ssm_client_stub, node_id, cmd_id)

    # given our mocks and the example CloudWatch Agent config as input...
    # expect CloudWatch Agent configured to use updated file on each cluster
    # node successfully
    cloudwatch_helper._update_cloudwatch_config("agent", is_head_node)

    # expect no pending responses left in client stub queues
    ec2_client_stub.assert_no_pending_responses()
    ssm_client_stub.assert_no_pending_responses()


def test_cloudwatch_agent_update_with_change_worker_node(
        ec2_client_stub, ssm_client_stub):
    # create test cluster worker node ID and an associated cloudwatch helper
    node_id = "i-abc"
    is_head_node = False
    cloudwatch_helper = helpers.get_cloudwatch_helper(node_id)

    # given a directive to check for the Unified CloudWatch Agent status...
    # expect CloudWatch Agent is installed
    stubs.get_ec2_cwa_installed_tag_true(ec2_client_stub, node_id)

    # given a directive to update a cluster CloudWatch Agent Config with new
    # changes...
    # expect the stored the CloudWatch Agent Config is different from local
    # config
    stubs.get_head_node_config_hash_different(ec2_client_stub, "agent",
                                              cloudwatch_helper, node_id)
    stubs.get_cur_node_config_hash_different(ec2_client_stub, "agent", node_id)

    # given an updated CloudWatch Agent Config file...
    # expect to update the node tag equal to updated config file sha1 hash
    # to reflect the changes in config file
    stubs.update_hash_tag_success(ec2_client_stub, node_id, "agent",
                                  cloudwatch_helper)
    # given that updated CloudWatch Agent Config is put to Parameter Store...
    # expect to send an SSM command to restart CloudWatch Agent on all nodes
    cmd_id = stubs.send_command_stop_cwa(ssm_client_stub, node_id)
    # given a SSM command to stop CloudWatch Agent sent to all nodes...
    # expect to wait for the command to complete successfully on every node
    stubs.list_command_invocations_success(ssm_client_stub, node_id, cmd_id)
    cw_ssm_param_name = helpers.get_ssm_param_name(
        cloudwatch_helper.cluster_name, "agent")
    cmd_id = stubs.send_command_start_cwa(ssm_client_stub, node_id,
                                          cw_ssm_param_name)
    # given a SSM command to start CloudWatch Agent sent to all nodes...
    # expect to wait for the command to complete successfully on every node
    stubs.list_command_invocations_success(ssm_client_stub, node_id, cmd_id)

    # given our mocks and the example CloudWatch Agent config as input...
    # expect CloudWatch Agent configured to use updated file on each cluster
    # node successfully
    cloudwatch_helper._update_cloudwatch_config("agent", is_head_node)

    # expect no pending responses left in client stub queues
    ec2_client_stub.assert_no_pending_responses()
    ssm_client_stub.assert_no_pending_responses()


def test_cloudwatch_dashboard_update_head_node(
        ec2_client_stub, ssm_client_stub, cloudwatch_client_stub):
    # create test cluster head node ID and an associated cloudwatch helper
    node_id = "i-abc"
    is_head_node = True
    cloudwatch_helper = helpers.get_cloudwatch_helper(node_id)

    # given a directive to check for the Unified CloudWatch Agent status...
    # expect CloudWatch Agent is installed
    stubs.get_ec2_cwa_installed_tag_true(ec2_client_stub, node_id)

    # given a directive to update a cluster CloudWatch Dashboard Config
    # with new changes...
    # expect the stored the CloudWatch Dashboard Config is different from local
    # config
    cw_ssm_param_name = helpers.get_ssm_param_name(
        cloudwatch_helper.cluster_name, "dashboard")
    stubs.get_param_ssm_different(ssm_client_stub, cw_ssm_param_name)

    # given an updated CloudWatch Dashboard Config file...
    # expect to store the new CloudWatch Dashboard config as an SSM parameter
    stubs.put_parameter_cloudwatch_config(
        ssm_client_stub, cloudwatch_helper.cluster_name, "dashboard")

    # given an updated CloudWatch Dashboard Config file...
    # expect to update the node tag equal to updated config file sha1 hash
    # to reflect the changes in config file
    stubs.update_hash_tag_success(ec2_client_stub, node_id, "dashboard",
                                  cloudwatch_helper)

    # given a directive to create a cluster CloudWatch dashboard...
    # expect to make a call to create a dashboard for each node in the cluster
    stubs.put_cluster_dashboard_success(
        cloudwatch_client_stub,
        cloudwatch_helper,
    )
    # given our mocks and the example CloudWatch Dashboard config as input...
    # expect CloudWatch Dashboard configured to use updated file
    # on each cluster node successfully
    cloudwatch_helper._update_cloudwatch_config("dashboard", is_head_node)

    # expect no pending responses left in client stub queues
    ec2_client_stub.assert_no_pending_responses()
    ssm_client_stub.assert_no_pending_responses()


def test_cloudwatch_dashboard_update_worker_node(
        ec2_client_stub, ssm_client_stub, cloudwatch_client_stub):
    # create test cluster worker node ID and an associated cloudwatch helper
    node_id = "i-abc"
    is_head_node = False
    cloudwatch_helper = helpers.get_cloudwatch_helper(node_id)

    # given a directive to check for the Unified CloudWatch Agent status...
    # expect CloudWatch Agent is installed
    stubs.get_ec2_cwa_installed_tag_true(ec2_client_stub, node_id)

    # given a directive to update a cluster CloudWatch Dashboard Config
    # with new changes...
    # expect the stored the CloudWatch Dashboard Config is different from local
    # config
    stubs.get_head_node_config_hash_different(ec2_client_stub, "dashboard",
                                              cloudwatch_helper, node_id)
    stubs.get_cur_node_config_hash_different(ec2_client_stub, "dashboard",
                                             node_id)

    # given an updated CloudWatch Dashboard Config file...
    # expect to update the node tag equal to updated config file sha1 hash
    # to reflect the changes in config file
    stubs.update_hash_tag_success(ec2_client_stub, node_id, "dashboard",
                                  cloudwatch_helper)

    # given our mocks and the example CloudWatch Dashboard config as input...
    # expect CloudWatch Dashboard configured to use updated file
    # on each cluster node successfully
    cloudwatch_helper._update_cloudwatch_config("dashboard", is_head_node)

    # expect no pending responses left in client stub queues
    ec2_client_stub.assert_no_pending_responses()
    ssm_client_stub.assert_no_pending_responses()


def test_cloudwatch_alarm_update_head_node(ec2_client_stub, ssm_client_stub,
                                           cloudwatch_client_stub):
    # create test cluster head node ID and an associated cloudwatch helper
    node_id = "i-abc"
    is_head_node = True
    cloudwatch_helper = helpers.get_cloudwatch_helper(node_id)

    # given a directive to check for the Unified CloudWatch Agent status...
    # expect CloudWatch Agent is installed
    stubs.get_ec2_cwa_installed_tag_true(ec2_client_stub, node_id)

    # given a directive to update a cluster CloudWatch Alarm Config with new
    # changes...
    # expect the stored the CloudWatch Alarm Config is different from local
    # config
    cw_ssm_param_name = helpers.get_ssm_param_name(
        cloudwatch_helper.cluster_name, "alarm")
    stubs.get_param_ssm_different(ssm_client_stub, cw_ssm_param_name)

    # given an updated CloudWatch Alarm Config file...
    # expect to store the new CloudWatch Alarm config as an SSM parameter
    stubs.put_parameter_cloudwatch_config(
        ssm_client_stub, cloudwatch_helper.cluster_name, "alarm")

    # given an updated CloudWatch Alarm Config file...
    # expect to update the node tag equal to updated config file sha1 hash
    # to reflect the changes in config file
    stubs.update_hash_tag_success(ec2_client_stub, node_id, "alarm",
                                  cloudwatch_helper)
    stubs.get_param_ssm_same(ssm_client_stub, cw_ssm_param_name,
                             cloudwatch_helper, "alarm")

    # given a directive to create cluster  CloudWatch Alarms...
    # expect to make a call to create alarms for each node in the cluster
    stubs.put_cluster_alarms_success(cloudwatch_client_stub, cloudwatch_helper)

    # given our mocks and the example  CloudWatch Alarm config as input...
    # expect  CloudWatch Alarm configured to use updated file on each cluster
    # node successfully
    cloudwatch_helper._update_cloudwatch_config("alarm", is_head_node)

    # expect no pending responses left in client stub queues
    ec2_client_stub.assert_no_pending_responses()
    ssm_client_stub.assert_no_pending_responses()


def test_cloudwatch_alarm_update_worker_node(ec2_client_stub, ssm_client_stub,
                                             cloudwatch_client_stub):
    # create test cluster worker node ID and an associated cloudwatch helper
    node_id = "i-abc"
    is_head_node = False
    cloudwatch_helper = helpers.get_cloudwatch_helper(node_id)

    # given a directive to check for the Unified CloudWatch Agent status...
    # expect CloudWatch Agent is installed
    stubs.get_ec2_cwa_installed_tag_true(ec2_client_stub, node_id)

    # given a directive to update a cluster CloudWatch Alarm Config with new
    # changes...
    # expect the stored the CloudWatch Alarm Config is different from local
    # config
    cw_ssm_param_name = helpers.get_ssm_param_name(
        cloudwatch_helper.cluster_name, "alarm")

    # given a directive to update a cluster CloudWatch Alarm Config with new
    # changes...
    # expect the stored the CloudWatch Alarm Config is different from local
    # config
    stubs.get_head_node_config_hash_different(ec2_client_stub, "alarm",
                                              cloudwatch_helper, node_id)
    stubs.get_cur_node_config_hash_different(ec2_client_stub, "alarm", node_id)

    # given an updated CloudWatch Alarm Config file...
    # expect to update the node tag equal to updated config file sha1 hash
    # to reflect the changes in config file
    stubs.update_hash_tag_success(ec2_client_stub, node_id, "alarm",
                                  cloudwatch_helper)
    stubs.get_param_ssm_same(ssm_client_stub, cw_ssm_param_name,
                             cloudwatch_helper, "alarm")

    # given a directive to create cluster CloudWatch Alarms...
    # expect to make a call to create alarms for each node in the cluster
    stubs.put_cluster_alarms_success(cloudwatch_client_stub, cloudwatch_helper)
    # given our mocks and the example CloudWatch Alarm config as input...
    # expect CloudWatch Alarm configured to use updated file on each cluster
    # node successfully
    cloudwatch_helper._update_cloudwatch_config("alarm", is_head_node)

    # expect no pending responses left in client stub queues
    ec2_client_stub.assert_no_pending_responses()
    ssm_client_stub.assert_no_pending_responses()


if __name__ == "__main__":
    import sys
    sys.exit(pytest.main(["-v", __file__]))