From 4944bb19fdc72a597d98d5d4ef39ee53920e4e1d Mon Sep 17 00:00:00 2001 From: Sylvain Hellegouarch Date: Tue, 26 Mar 2024 09:49:13 +0100 Subject: [PATCH] add set_tags_on_instances and remove_tags_from_instances Signed-off-by: Sylvain Hellegouarch --- CHANGELOG.md | 5 ++ chaosaws/ec2/actions.py | 112 +++++++++++++++++++++++++++++++++- tests/conftest.py | 1 - tests/ec2/test_ec2_actions.py | 52 ++++++++++++++++ 4 files changed, 167 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4fcd2fd..9fcf80e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,11 @@ [Unreleased]: https://github.com/chaostoolkit-incubator/chaostoolkit-aws/compare/0.32.1...HEAD +### Added + +* The `set_tags_on_instances` and `remove_tags_from_instances` actions on the + EC2 package + ## [0.32.1][] - 2024-02-23 [0.32.1]: https://github.com/chaostoolkit-incubator/chaostoolkit-aws/compare/0.32.0...0.32.1 diff --git a/chaosaws/ec2/actions.py b/chaosaws/ec2/actions.py index 99f70c1..2581e1e 100644 --- a/chaosaws/ec2/actions.py +++ b/chaosaws/ec2/actions.py @@ -9,7 +9,12 @@ from chaoslib.exceptions import ActivityFailed, FailedActivity from chaoslib.types import Configuration, Secrets -from chaosaws import aws_client, convert_tags, get_logger +from chaosaws import ( + aws_client, + convert_tags, + get_logger, + tags_as_key_value_pairs, +) from chaosaws.types import AWSResponse __all__ = [ @@ -22,6 +27,8 @@ "detach_random_volume", "attach_volume", "stop_instances_by_incremental_steps", + "set_tags_on_instances", + "remove_tags_from_instances", ] logger = get_logger() @@ -496,6 +503,107 @@ def stop_instances_by_incremental_steps( return responses +def set_tags_on_instances( + tags: Union[str, List[Dict[str, str]]], + percentage: int = 100, + az: str = None, + filters: List[Dict[str, Any]] = None, + configuration: Configuration = None, + secrets: Secrets = None, +) -> AWSResponse: + """ + Sets some tags on the instances matching the `filters`. The set of instances + may be filtered down by availability-zone too. + + The `tags`can be passed as a dictionary of key, value pair respecting + the usual AWS form: [{"Key": "...", "Value": "..."}, ...] or as a string + of key value pairs such as "k1=v1,k2=v2" + + The `percentage` parameter (between 0 and 100) allows you to select only a + certain amount of instances amongst those matching the filters. + + If no filters are given and `percentage` remains to 100, the entire set + of instances in an AZ will be tagged. If no AZ is provided, your entire + set of instances in the region will be tagged. This can be a lot of + instances and would not be appropriate. Always to use the filters to + target a significant subset. + + See also: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ec2/client/create_tags.html + """ # noqa E501 + client = aws_client("ec2", configuration, secrets) + + if isinstance(tags, str): + tags = tags_as_key_value_pairs(convert_tags(tags) if tags else []) + + if not tags: + raise FailedActivity("Missing tags to be set") + + filters = filters or [] + if az: + filters.append({"Name": "availability-zone", "Values": [az]}) + + instances = list_instances_by_type(filters, client) + + instance_ids = [inst_id for inst_id in instances.get("normal", [])] + + total = len(instance_ids) + # always force at least one instance + count = max(1, round(total * percentage / 100)) + target_instances = random.sample(instance_ids, count) + + if not target_instances: + raise FailedActivity(f"No instances in availability zone: {az}") + + logger.debug( + "Picked EC2 instances '{}' from AZ '{}'".format( + str(target_instances), az + ) + ) + + response = client.create_tags(Resources=target_instances, Tags=tags) + + return response + + +def remove_tags_from_instances( + tags: Union[str, List[Dict[str, str]]], + az: str = None, + configuration: Configuration = None, + secrets: Secrets = None, +) -> AWSResponse: + """ + Remove tags from instances + + Usually mirrors `set_tags_on_instances`. + """ + client = aws_client("ec2", configuration, secrets) + + if isinstance(tags, str): + tags = tags_as_key_value_pairs(convert_tags(tags) if tags else []) + + filters = [] + for tag in tags: + filters.append({"Name": f"tag:{tag['Key']}", "Values": [tag["Value"]]}) + + if az: + filters.append({"Name": "availability-zone", "Values": [az]}) + + instances = client.describe_instances(Filters=filters) + + instance_ids = [] + for reservation in instances["Reservations"]: + for inst in reservation["Instances"]: + instance_ids.append(inst["InstanceId"]) + + logger.debug( + "Found EC2 instances '{}' from AZ '{}'".format(str(instance_ids), az) + ) + + response = client.delete_tags(Resources=instance_ids, Tags=tags) + + return response + + ############################################################################### # Private functions ############################################################################### @@ -604,7 +712,7 @@ def get_instance_type_from_response(response: Dict) -> Dict: """ Transform list of instance IDs to a dict with IDs by instance type """ - instances_type = defaultdict(List) + instances_type = defaultdict(list) # reservations are instances that were started together for reservation in response["Reservations"]: diff --git a/tests/conftest.py b/tests/conftest.py index 5578f8b..f0d4c2a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -12,5 +12,4 @@ @pytest.fixture(scope="session", autouse=True) def setup_logger() -> None: - print("#######################") configure_logger(verbose=True) diff --git a/tests/ec2/test_ec2_actions.py b/tests/ec2/test_ec2_actions.py index bf9b758..db79962 100644 --- a/tests/ec2/test_ec2_actions.py +++ b/tests/ec2/test_ec2_actions.py @@ -8,8 +8,10 @@ attach_volume, authorize_security_group_ingress, detach_random_volume, + remove_tags_from_instances, restart_instances, revoke_security_group_ingress, + set_tags_on_instances, start_instances, stop_instance, stop_instances, @@ -1134,3 +1136,53 @@ def test_revoke_security_group_ingress_with_cidr_ip(aws_client): } ], ) + + +@patch("chaosaws.ec2.actions.aws_client", autospec=True) +def test_set_tags_on_instances(aws_client): + tags = "a=b,c=d" + expected_tags = [{"Key": "a", "Value": "b"}, {"Key": "c", "Value": "d"}] + + client = MagicMock() + aws_client.return_value = client + inst_id_1 = "i-1234567890abcdef0" + client.describe_instances.return_value = { + "Reservations": [ + { + "Instances": [ + {"InstanceId": inst_id_1, "InstanceLifecycle": "normal"} + ] + } + ] + } + + set_tags_on_instances(tags, percentage=10) + + client.create_tags.assert_called_with( + Resources=[inst_id_1], Tags=expected_tags + ) + + +@patch("chaosaws.ec2.actions.aws_client", autospec=True) +def test_remove_tags_from_instances(aws_client): + tags = "a=b,c=d" + expected_tags = [{"Key": "a", "Value": "b"}, {"Key": "c", "Value": "d"}] + + client = MagicMock() + aws_client.return_value = client + inst_id_1 = "i-1234567890abcdef0" + client.describe_instances.return_value = { + "Reservations": [ + { + "Instances": [ + {"InstanceId": inst_id_1, "InstanceLifecycle": "normal"} + ] + } + ] + } + + remove_tags_from_instances(tags) + + client.delete_tags.assert_called_with( + Resources=[inst_id_1], Tags=expected_tags + )