Add Tag support for other resource types then an account (#3982)

* - Adding checking for resource type in tag functions
- Adding TargetNotFoundException when no resource found
- Adding support for tags for root OU, OU and Policies
- Adding tests covering the new code
- Adding test for deletion of a tag

* fixed linting issue

* - renamed helper function to a more logical name
- added tests for helper function
- fixed bugs in tests for tag functions

Co-authored-by: Sjoerd Tromp <stromp@schubergphilis.com>

moto/organizations/models.py

@@ -71,7 +71,7 @@ class FakeAccount(BaseModel):
         self.joined_method = "CREATED"
         self.parent_id = organization.root_id
         self.attached_policies = []
-        self.tags = {}
+        self.tags = {tag["Key"]: tag["Value"] for tag in kwargs.get("Tags", [])}
 
     @property
     def arn(self):
@@ -114,6 +114,7 @@ class FakeOrganizationalUnit(BaseModel):
         self.parent_id = kwargs.get("ParentId")
         self._arn_format = utils.OU_ARN_FORMAT
         self.attached_policies = []
+        self.tags = {tag["Key"]: tag["Value"] for tag in kwargs.get("Tags", [])}
 
     @property
     def arn(self):
@@ -143,6 +144,7 @@ class FakeRoot(FakeOrganizationalUnit):
         self.policy_types = [{"Type": "SERVICE_CONTROL_POLICY", "Status": "ENABLED"}]
         self._arn_format = utils.ROOT_ARN_FORMAT
         self.attached_policies = []
+        self.tags = {tag["Key"]: tag["Value"] for tag in kwargs.get("Tags", [])}
 
     def describe(self):
         return {
@@ -654,6 +656,25 @@ class OrganizationsBackend(BaseBackend):
             ]
         )
 
+    def _get_resource_for_tagging(self, resource_id):
+        if re.compile(utils.OU_ID_REGEX).fullmatch(resource_id) or re.fullmatch(
+            utils.ROOT_ID_REGEX, resource_id
+        ):
+            resource = next((a for a in self.ou if a.id == resource_id), None)
+        elif re.compile(utils.ACCOUNT_ID_REGEX).fullmatch(resource_id):
+            resource = next((a for a in self.accounts if a.id == resource_id), None)
+        elif re.compile(utils.POLICY_ID_REGEX).fullmatch(resource_id):
+            resource = next((a for a in self.policies if a.id == resource_id), None)
+        else:
+            raise InvalidInputException(
+                "You provided a value that does not match the required pattern."
+            )
+
+        if resource is None:
+            raise TargetNotFoundException
+
+        return resource
+
     def list_targets_for_policy(self, **kwargs):
         if re.compile(utils.POLICY_ID_REGEX).match(kwargs["PolicyId"]):
             policy = next(
@@ -673,37 +694,19 @@ class OrganizationsBackend(BaseBackend):
         return dict(Targets=objects)
 
     def tag_resource(self, **kwargs):
-        account = next((a for a in self.accounts if a.id == kwargs["ResourceId"]), None)
+        resource = self._get_resource_for_tagging(kwargs["ResourceId"])
-
-        if account is None:
-            raise InvalidInputException(
-                "You provided a value that does not match the required pattern."
-            )
-
         new_tags = {tag["Key"]: tag["Value"] for tag in kwargs["Tags"]}
-        account.tags.update(new_tags)
+        resource.tags.update(new_tags)
 
     def list_tags_for_resource(self, **kwargs):
-        account = next((a for a in self.accounts if a.id == kwargs["ResourceId"]), None)
+        resource = self._get_resource_for_tagging(kwargs["ResourceId"])
-
+        tags = [{"Key": key, "Value": value} for key, value in resource.tags.items()]
-        if account is None:
-            raise InvalidInputException(
-                "You provided a value that does not match the required pattern."
-            )
-
-        tags = [{"Key": key, "Value": value} for key, value in account.tags.items()]
         return dict(Tags=tags)
 
     def untag_resource(self, **kwargs):
-        account = next((a for a in self.accounts if a.id == kwargs["ResourceId"]), None)
+        resource = self._get_resource_for_tagging(kwargs["ResourceId"])
-
-        if account is None:
-            raise InvalidInputException(
-                "You provided a value that does not match the required pattern."
-            )
-
         for key in kwargs["TagKeys"]:
-            account.tags.pop(key, None)
+            resource.tags.pop(key, None)
 
     def enable_aws_service_access(self, **kwargs):
         service = FakeServiceAccess(**kwargs)

tests/test_organizations/test_organizations_boto3.py

@@ -2,6 +2,16 @@ from __future__ import unicode_literals
 
 from datetime import datetime
 
+from moto.organizations.exceptions import InvalidInputException, TargetNotFoundException
+from moto.organizations.models import (
+    FakeAccount,
+    FakeOrganization,
+    FakeOrganizationalUnit,
+    FakePolicy,
+    FakeRoot,
+    OrganizationsBackend,
+)
+
 import boto3
 import json
 import six
@@ -966,26 +976,84 @@ def test_list_targets_for_policy_exception():
 
 
 @mock_organizations
-def test_tag_resource():
+def test_tag_resource_account():
     client = boto3.client("organizations", region_name="us-east-1")
     client.create_organization(FeatureSet="ALL")
-    account_id = client.create_account(AccountName=mockname, Email=mockemail)[
+    resource_id = client.create_account(AccountName=mockname, Email=mockemail)[
         "CreateAccountStatus"
     ]["AccountId"]
 
-    client.tag_resource(ResourceId=account_id, Tags=[{"Key": "key", "Value": "value"}])
+    client.tag_resource(ResourceId=resource_id, Tags=[{"Key": "key", "Value": "value"}])
 
-    response = client.list_tags_for_resource(ResourceId=account_id)
+    response = client.list_tags_for_resource(ResourceId=resource_id)
     response["Tags"].should.equal([{"Key": "key", "Value": "value"}])
 
     # adding a tag with an existing key, will update the value
     client.tag_resource(
-        ResourceId=account_id, Tags=[{"Key": "key", "Value": "new-value"}]
+        ResourceId=resource_id, Tags=[{"Key": "key", "Value": "new-value"}]
     )
 
-    response = client.list_tags_for_resource(ResourceId=account_id)
+    response = client.list_tags_for_resource(ResourceId=resource_id)
     response["Tags"].should.equal([{"Key": "key", "Value": "new-value"}])
 
+    client.untag_resource(ResourceId=resource_id, TagKeys=["key"])
+
+    response = client.list_tags_for_resource(ResourceId=resource_id)
+    response["Tags"].should.equal([])
+
+
+@mock_organizations
+def test_tag_resource_organization_organization_root():
+    client = boto3.client("organizations", region_name="us-east-1")
+    client.create_organization(FeatureSet="ALL")
+
+    resource_id = client.list_roots()["Roots"][0]["Id"]
+    client.tag_resource(ResourceId=resource_id, Tags=[{"Key": "key", "Value": "value"}])
+
+    response = client.list_tags_for_resource(ResourceId=resource_id)
+    response["Tags"].should.equal([{"Key": "key", "Value": "value"}])
+
+    # adding a tag with an existing key, will update the value
+    client.tag_resource(
+        ResourceId=resource_id, Tags=[{"Key": "key", "Value": "new-value"}]
+    )
+
+    response = client.list_tags_for_resource(ResourceId=resource_id)
+    response["Tags"].should.equal([{"Key": "key", "Value": "new-value"}])
+
+    client.untag_resource(ResourceId=resource_id, TagKeys=["key"])
+
+    response = client.list_tags_for_resource(ResourceId=resource_id)
+    response["Tags"].should.equal([])
+
+
+@mock_organizations
+def test_tag_resource_organization_organizational_unit():
+    client = boto3.client("organizations", region_name="us-east-1")
+    client.create_organization(FeatureSet="ALL")
+    root_id = client.list_roots()["Roots"][0]["Id"]
+    resource_id = client.create_organizational_unit(ParentId=root_id, Name="ou01")[
+        "OrganizationalUnit"
+    ]["Id"]
+
+    client.tag_resource(ResourceId=resource_id, Tags=[{"Key": "key", "Value": "value"}])
+
+    response = client.list_tags_for_resource(ResourceId=resource_id)
+    response["Tags"].should.equal([{"Key": "key", "Value": "value"}])
+
+    # adding a tag with an existing key, will update the value
+    client.tag_resource(
+        ResourceId=resource_id, Tags=[{"Key": "key", "Value": "new-value"}]
+    )
+
+    response = client.list_tags_for_resource(ResourceId=resource_id)
+    response["Tags"].should.equal([{"Key": "key", "Value": "new-value"}])
+
+    client.untag_resource(ResourceId=resource_id, TagKeys=["key"])
+
+    response = client.list_tags_for_resource(ResourceId=resource_id)
+    response["Tags"].should.equal([])
+
 
 @mock_organizations
 def test_tag_resource_errors():
@@ -994,7 +1062,7 @@ def test_tag_resource_errors():
 
     with pytest.raises(ClientError) as e:
         client.tag_resource(
-            ResourceId="000000000000", Tags=[{"Key": "key", "Value": "value"},],
+            ResourceId="0A000000X000", Tags=[{"Key": "key", "Value": "value"},],
         )
     ex = e.value
     ex.operation_name.should.equal("TagResource")
@@ -1003,6 +1071,119 @@ def test_tag_resource_errors():
     ex.response["Error"]["Message"].should.equal(
         "You provided a value that does not match the required pattern."
     )
+    with pytest.raises(ClientError) as e:
+        client.tag_resource(
+            ResourceId="000000000000", Tags=[{"Key": "key", "Value": "value"}]
+        )
+    ex = e.value
+    ex.operation_name.should.equal("TagResource")
+    ex.response["ResponseMetadata"]["HTTPStatusCode"].should.equal(400)
+    ex.response["Error"]["Code"].should.contain("TargetNotFoundException")
+    ex.response["Error"]["Message"].should.equal(
+        "You specified a target that doesn't exist."
+    )
+
+
+def test__get_resource_for_tagging_existing_root():
+    org = FakeOrganization("ALL")
+    root = FakeRoot(org)
+
+    org_backend = OrganizationsBackend()
+    org_backend.ou.append(root)
+    response = org_backend._get_resource_for_tagging(root.id)
+    response.id.should.equal(root.id)
+
+
+def test__get_resource_for_tagging_existing_non_root():
+    org_backend = OrganizationsBackend()
+    with pytest.raises(TargetNotFoundException) as e:
+        org_backend._get_resource_for_tagging("r-abcd")
+    ex = e.value
+    ex.code.should.equal(400)
+    ex.description.should.contain("TargetNotFoundException")
+    ex.message.should.equal("You specified a target that doesn't exist.")
+
+
+def test__get_resource_for_tagging_existing_ou():
+    org = FakeOrganization("ALL")
+    ou = FakeOrganizationalUnit(org)
+    org_backend = OrganizationsBackend()
+
+    org_backend.ou.append(ou)
+    response = org_backend._get_resource_for_tagging(ou.id)
+    response.id.should.equal(ou.id)
+
+
+def test__get_resource_for_tagging_non_existing_ou():
+    org_backend = OrganizationsBackend()
+    with pytest.raises(TargetNotFoundException) as e:
+        org_backend._get_resource_for_tagging("ou-9oyc-lv2q36ln")
+    ex = e.value
+    ex.code.should.equal(400)
+    ex.description.should.contain("TargetNotFoundException")
+    ex.message.should.equal("You specified a target that doesn't exist.")
+
+
+def test__get_resource_for_tagging_existing_account():
+    org = FakeOrganization("ALL")
+    org_backend = OrganizationsBackend()
+    account = FakeAccount(org, AccountName="test", Email="test@test.test")
+
+    org_backend.accounts.append(account)
+    response = org_backend._get_resource_for_tagging(account.id)
+    response.id.should.equal(account.id)
+
+
+def test__get_resource_for_tagging_non_existing_account():
+    org_backend = OrganizationsBackend()
+    with pytest.raises(TargetNotFoundException) as e:
+        org_backend._get_resource_for_tagging("100326223992")
+    ex = e.value
+    ex.code.should.equal(400)
+    ex.description.should.contain("TargetNotFoundException")
+    ex.message.should.equal("You specified a target that doesn't exist.")
+
+
+def test__get_resource_for_tagging_existing_policy():
+    org = FakeOrganization("ALL")
+    org_backend = OrganizationsBackend()
+    policy = FakePolicy(org, Type="SERVICE_CONTROL_POLICY")
+
+    org_backend.policies.append(policy)
+    response = org_backend._get_resource_for_tagging(policy.id)
+    response.id.should.equal(policy.id)
+
+
+def test__get_resource_for_tagging_non_existing_policy():
+    org_backend = OrganizationsBackend()
+    with pytest.raises(TargetNotFoundException) as e:
+        org_backend._get_resource_for_tagging("p-y1vas4da")
+    ex = e.value
+    ex.code.should.equal(400)
+    ex.description.should.contain("TargetNotFoundException")
+    ex.message.should.equal("You specified a target that doesn't exist.")
+
+
+def test__get_resource_for_tagging_non_existing_policy():
+    org_backend = OrganizationsBackend()
+    with pytest.raises(TargetNotFoundException) as e:
+        org_backend._get_resource_for_tagging("p-y1vas4da")
+    ex = e.value
+    ex.code.should.equal(400)
+    ex.description.should.contain("TargetNotFoundException")
+    ex.message.should.equal("You specified a target that doesn't exist.")
+
+
+def test__get_resource_to_tag_incorrect_resource():
+    org_backend = OrganizationsBackend()
+    with pytest.raises(InvalidInputException) as e:
+        org_backend._get_resource_for_tagging("10032622399200")
+    ex = e.value
+    ex.code.should.equal(400)
+    ex.description.should.contain("InvalidInputException")
+    ex.message.should.equal(
+        "You provided a value that does not match the required pattern."
+    )
 
 
 @mock_organizations
@@ -1025,7 +1206,7 @@ def test_list_tags_for_resource_errors():
     client.create_organization(FeatureSet="ALL")
 
     with pytest.raises(ClientError) as e:
-        client.list_tags_for_resource(ResourceId="000000000000")
+        client.list_tags_for_resource(ResourceId="000x00000A00")
     ex = e.value
     ex.operation_name.should.equal("ListTagsForResource")
     ex.response["ResponseMetadata"]["HTTPStatusCode"].should.equal(400)
@@ -1033,6 +1214,15 @@ def test_list_tags_for_resource_errors():
     ex.response["Error"]["Message"].should.equal(
         "You provided a value that does not match the required pattern."
     )
+    with pytest.raises(ClientError) as e:
+        client.list_tags_for_resource(ResourceId="000000000000")
+    ex = e.value
+    ex.operation_name.should.equal("ListTagsForResource")
+    ex.response["ResponseMetadata"]["HTTPStatusCode"].should.equal(400)
+    ex.response["Error"]["Code"].should.contain("TargetNotFoundException")
+    ex.response["Error"]["Message"].should.equal(
+        "You specified a target that doesn't exist."
+    )
 
 
 @mock_organizations
@@ -1062,7 +1252,7 @@ def test_untag_resource_errors():
     client.create_organization(FeatureSet="ALL")
 
     with pytest.raises(ClientError) as e:
-        client.untag_resource(ResourceId="000000000000", TagKeys=["key"])
+        client.untag_resource(ResourceId="0X00000000A0", TagKeys=["key"])
     ex = e.value
     ex.operation_name.should.equal("UntagResource")
     ex.response["ResponseMetadata"]["HTTPStatusCode"].should.equal(400)
@@ -1070,6 +1260,15 @@ def test_untag_resource_errors():
     ex.response["Error"]["Message"].should.equal(
         "You provided a value that does not match the required pattern."
     )
+    with pytest.raises(ClientError) as e:
+        client.untag_resource(ResourceId="000000000000", TagKeys=["key"])
+    ex = e.value
+    ex.operation_name.should.equal("UntagResource")
+    ex.response["ResponseMetadata"]["HTTPStatusCode"].should.equal(400)
+    ex.response["Error"]["Code"].should.contain("TargetNotFoundException")
+    ex.response["Error"]["Message"].should.equal(
+        "You specified a target that doesn't exist."
+    )
 
 
 @mock_organizations