diff --git a/moto/elbv2/models.py b/moto/elbv2/models.py index 6bc8f860e..3c6afe7f5 100644 --- a/moto/elbv2/models.py +++ b/moto/elbv2/models.py @@ -65,6 +65,7 @@ class FakeTargetGroup(BaseModel): self.healthy_threshold_count = healthy_threshold_count self.unhealthy_threshold_count = unhealthy_threshold_count self.load_balancer_arns = [] + self.tags = {} self.attributes = { 'deregistration_delay.timeout_seconds': 300, @@ -86,6 +87,11 @@ class FakeTargetGroup(BaseModel): if not t: raise InvalidTargetError() + def add_tag(self, key, value): + if len(self.tags) >= 10 and key not in self.tags: + raise TooManyTagsError() + self.tags[key] = value + def health_for(self, target): t = self.targets.get(target['id']) if t is None: diff --git a/moto/elbv2/responses.py b/moto/elbv2/responses.py index 652ecc566..3e8535187 100644 --- a/moto/elbv2/responses.py +++ b/moto/elbv2/responses.py @@ -266,10 +266,17 @@ class ELBV2Response(BaseResponse): resource_arns = self._get_multi_param('ResourceArns.member') for arn in resource_arns: - load_balancer = self.elbv2_backend.load_balancers.get(arn) - if not load_balancer: + if ':targetgroup' in arn: + resource = self.elbv2_backend.target_groups.get(arn) + if not resource: + raise TargetGroupNotFoundError() + elif ':loadbalancer' in arn: + resource = self.elbv2_backend.load_balancers.get(arn) + if not resource: + raise LoadBalancerNotFoundError() + else: raise LoadBalancerNotFoundError() - self._add_tags(load_balancer) + self._add_tags(resource) template = self.response_template(ADD_TAGS_TEMPLATE) return template.render() @@ -279,30 +286,41 @@ class ELBV2Response(BaseResponse): tag_keys = self._get_multi_param('TagKeys.member') for arn in resource_arns: - load_balancer = self.elbv2_backend.load_balancers.get(arn) - if not load_balancer: + if ':targetgroup' in arn: + resource = self.elbv2_backend.target_groups.get(arn) + if not resource: + raise TargetGroupNotFoundError() + elif ':loadbalancer' in arn: + resource = self.elbv2_backend.load_balancers.get(arn) + if not resource: + raise LoadBalancerNotFoundError() + else: raise LoadBalancerNotFoundError() - [load_balancer.remove_tag(key) for key in tag_keys] + [resource.remove_tag(key) for key in tag_keys] template = self.response_template(REMOVE_TAGS_TEMPLATE) return template.render() def describe_tags(self): - elbs = [] - for key, value in self.querystring.items(): - if "ResourceArns.member" in key: - number = key.split('.')[2] - load_balancer_arn = self._get_param( - 'ResourceArns.member.{0}'.format(number)) - elb = self.elbv2_backend.load_balancers.get(load_balancer_arn) - if not elb: + resource_arns = self._get_multi_param('ResourceArns.member') + resources = [] + for arn in resource_arns: + if ':targetgroup' in arn: + resource = self.elbv2_backend.target_groups.get(arn) + if not resource: + raise TargetGroupNotFoundError() + elif ':loadbalancer' in arn: + resource = self.elbv2_backend.load_balancers.get(arn) + if not resource: raise LoadBalancerNotFoundError() - elbs.append(elb) + else: + raise LoadBalancerNotFoundError() + resources.append(resource) template = self.response_template(DESCRIBE_TAGS_TEMPLATE) - return template.render(load_balancers=elbs) + return template.render(resources=resources) - def _add_tags(self, elb): + def _add_tags(self, resource): tag_values = [] tag_keys = [] @@ -324,7 +342,7 @@ class ELBV2Response(BaseResponse): raise DuplicateTagKeysError(counts[0]) for tag_key, tag_value in zip(tag_keys, tag_values): - elb.add_tag(tag_key, tag_value) + resource.add_tag(tag_key, tag_value) ADD_TAGS_TEMPLATE = """ @@ -344,11 +362,11 @@ REMOVE_TAGS_TEMPLATE = """ - {% for load_balancer in load_balancers %} + {% for resource in resources %} - {{ load_balancer.arn }} + {{ resource.arn }} - {% for key, value in load_balancer.tags.items() %} + {% for key, value in resource.tags.items() %} {{ value }} {{ key }} diff --git a/tests/test_elbv2/test_elbv2.py b/tests/test_elbv2/test_elbv2.py index e84cd0080..21799ddcf 100644 --- a/tests/test_elbv2/test_elbv2.py +++ b/tests/test_elbv2/test_elbv2.py @@ -251,6 +251,12 @@ def test_create_target_group_and_listeners(): UnhealthyThresholdCount=2, Matcher={'HttpCode': '200'}) target_group = response.get('TargetGroups')[0] + target_group_arn = target_group['TargetGroupArn'] + + # Add tags to the target group + conn.add_tags(ResourceArns=[target_group_arn], Tags=[{'Key': 'target', 'Value': 'group'}]) + conn.describe_tags(ResourceArns=[target_group_arn])['TagDescriptions'][0]['Tags'].should.equal( + [{'Key': 'target', 'Value': 'group'}]) # Check it's in the describe_target_groups response response = conn.describe_target_groups()