wehub-resource-sync
1190 lines
48 KiB
Python
1190 lines
48 KiB
Python
import copy
|
|
import sys
|
|
from unittest.mock import Mock, patch
|
|
|
|
import pytest
|
|
from click.exceptions import ClickException
|
|
|
|
import ray.tests.aws.utils.helpers as helpers
|
|
import ray.tests.aws.utils.stubs as stubs
|
|
from ray.autoscaler._private.aws.config import (
|
|
DEFAULT_AMI,
|
|
_configure_subnet,
|
|
_get_subnets_or_die,
|
|
bootstrap_aws,
|
|
log_to_cli,
|
|
)
|
|
from ray.autoscaler._private.aws.node_provider import AWSNodeProvider
|
|
from ray.autoscaler._private.providers import _get_node_provider
|
|
from ray.tests.aws.utils.constants import (
|
|
AUX_SG,
|
|
AUX_SUBNET,
|
|
CUSTOM_IN_BOUND_RULES,
|
|
DEFAULT_CLUSTER_NAME,
|
|
DEFAULT_INSTANCE_PROFILE,
|
|
DEFAULT_KEY_PAIR,
|
|
DEFAULT_LT,
|
|
DEFAULT_SG,
|
|
DEFAULT_SG_AUX_SUBNET,
|
|
DEFAULT_SG_DUAL_GROUP_RULES,
|
|
DEFAULT_SG_WITH_NAME,
|
|
DEFAULT_SG_WITH_NAME_AND_RULES,
|
|
DEFAULT_SG_WITH_RULES,
|
|
DEFAULT_SG_WITH_RULES_AUX_SUBNET,
|
|
DEFAULT_SUBNET,
|
|
)
|
|
|
|
|
|
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_subnets_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"]]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"correct_az",
|
|
[True, False],
|
|
)
|
|
def test_create_sg_different_vpc_same_rules(
|
|
iam_client_stub, ec2_client_stub, correct_az: bool
|
|
):
|
|
# use default stubs to skip ahead to security group configuration
|
|
stubs.skip_to_configure_sg(ec2_client_stub, iam_client_stub)
|
|
|
|
default_subnet = copy.deepcopy(DEFAULT_SUBNET)
|
|
if not correct_az:
|
|
default_subnet["AvailabilityZone"] = "us-west-2b"
|
|
|
|
# given head and worker nodes with custom subnets defined...
|
|
# expect to second describe the head subnet ID
|
|
stubs.describe_subnets_echo(ec2_client_stub, [default_subnet])
|
|
# expect to first describe the worker subnet ID
|
|
stubs.describe_subnets_echo(ec2_client_stub, [AUX_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
|
|
error = None
|
|
try:
|
|
config = helpers.bootstrap_aws_example_config_file("example-subnets.yaml")
|
|
except ClickException as e:
|
|
error = e
|
|
|
|
_get_subnets_or_die.cache_clear()
|
|
|
|
if not correct_az:
|
|
assert isinstance(error, ClickException), "Did not get a ClickException!"
|
|
iam_client_stub._queue.clear()
|
|
ec2_client_stub._queue.clear()
|
|
return
|
|
|
|
# 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_subnets_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()
|
|
|
|
|
|
# Parametrize across multiple regions, since default AMI is different in each
|
|
@pytest.mark.parametrize(
|
|
"iam_client_stub,ec2_client_stub,region",
|
|
[3 * (region,) for region in DEFAULT_AMI],
|
|
indirect=["iam_client_stub", "ec2_client_stub"],
|
|
)
|
|
def test_fills_out_amis_and_iam(iam_client_stub, ec2_client_stub, region):
|
|
# Set up expected key pair for specific region
|
|
region_key_pair = DEFAULT_KEY_PAIR.copy()
|
|
region_key_pair["KeyName"] = DEFAULT_KEY_PAIR["KeyName"].replace(
|
|
"us-west-2", region
|
|
)
|
|
|
|
# Setup stubs to mock out boto3
|
|
stubs.configure_iam_role_default(iam_client_stub)
|
|
stubs.configure_key_pair_default(
|
|
ec2_client_stub, region=region, expected_key_pair=region_key_pair
|
|
)
|
|
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"]
|
|
|
|
config["provider"]["region"] = region
|
|
|
|
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_subnets_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)
|
|
stubs.describe_subnets_echo(
|
|
ec2_client_stub,
|
|
[DEFAULT_SUBNET, {**DEFAULT_SUBNET, "SubnetId": "subnet-11111111"}],
|
|
)
|
|
stubs.describe_subnets_echo(
|
|
ec2_client_stub, [{**DEFAULT_SUBNET, "SubnetId": "subnet-22222222"}]
|
|
)
|
|
stubs.describe_subnets_echo(
|
|
ec2_client_stub, [{**DEFAULT_SUBNET, "SubnetId": "subnet-33333333"}]
|
|
)
|
|
|
|
# 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_resource"):
|
|
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_retry_get_node():
|
|
"""
|
|
This tests _get_node() retries `ec2.instances.filter` if the first call returns an empty list.
|
|
This is important because the EC2 API is eventually consistent, and it may take some time for the
|
|
instance to be available after it has been launched.
|
|
"""
|
|
|
|
with patch("ray.autoscaler._private.aws.node_provider.make_ec2_resource"):
|
|
provider = AWSNodeProvider(
|
|
provider_config={"region": "nowhere"},
|
|
cluster_name="default",
|
|
)
|
|
|
|
attempts = 0
|
|
fake_instance = {"id": "i-1234567890abcdef0"}
|
|
|
|
def mock_filter(*args, **kwargs):
|
|
nonlocal attempts
|
|
if kwargs.get("InstanceIds") == [fake_instance["id"]]:
|
|
attempts += 1
|
|
if attempts > 1:
|
|
return [fake_instance]
|
|
return []
|
|
|
|
provider.ec2.instances.filter.side_effect = mock_filter
|
|
|
|
assert provider._get_node(fake_instance["id"]) == fake_instance
|
|
assert attempts == 2
|
|
|
|
|
|
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__":
|
|
sys.exit(pytest.main(["-v", __file__]))
|