Add tests for KMS resource tagging and untagging (#12121)

localstack · Jan 20, 2025 · 7e62817 · 7e62817
1 parent bd8bfa6
commit 7e62817
Show file tree

Hide file tree

Showing 3 changed files with 211 additions and 41 deletions.
diff --git a/tests/aws/services/kms/test_kms.py b/tests/aws/services/kms/test_kms.py
@@ -21,6 +21,10 @@
 from localstack.utils.strings import short_uid, to_str
 
 
+def create_tags(**kwargs):
+    return [{"TagKey": key, "TagValue": value} for key, value in kwargs.items()]
+
+
 @pytest.fixture(scope="class")
 def kms_client_for_region(aws_client_factory):
     def _kms_client(
@@ -103,6 +107,99 @@ def test_create_key(
         assert f":{region_name}:" in response["Arn"]
         assert f":{account_id}:" in response["Arn"]
 
+    @markers.aws.validated
+    def test_tag_existing_key_and_untag(
+        self, kms_client_for_region, kms_create_key, snapshot, region_name
+    ):
+        kms_client = kms_client_for_region(region_name)
+        key_id = kms_create_key(
+            region_name=region_name, Description="test key 123", KeyUsage="ENCRYPT_DECRYPT"
+        )["KeyId"]
+
+        tags = create_tags(tag1="value1", tag2="value2")
+        kms_client.tag_resource(KeyId=key_id, Tags=tags)
+
+        response = kms_client.list_resource_tags(KeyId=key_id)["Tags"]
+        snapshot.match("list-resource-tags", response)
+
+        tag_keys = [tag["TagKey"] for tag in tags]
+        kms_client.untag_resource(KeyId=key_id, TagKeys=tag_keys)
+
+        response = kms_client.list_resource_tags(KeyId=key_id)["Tags"]
+        snapshot.match("list-resource-tags-after-all-untagged", response)
+
+    @markers.aws.validated
+    def test_create_key_with_tag_and_untag(
+        self, kms_client_for_region, kms_create_key, snapshot, region_name
+    ):
+        kms_client = kms_client_for_region(region_name)
+
+        tags = create_tags(tag1="value1", tag2="value2")
+        key_id = kms_create_key(
+            region_name=region_name,
+            Description="test key 123",
+            KeyUsage="ENCRYPT_DECRYPT",
+            Tags=tags,
+        )["KeyId"]
+
+        response = kms_client.list_resource_tags(KeyId=key_id)["Tags"]
+        snapshot.match("list-resource-tags", response)
+
+        tag_keys = [tag["TagKey"] for tag in tags]
+        kms_client.untag_resource(KeyId=key_id, TagKeys=tag_keys)
+
+        response = kms_client.list_resource_tags(KeyId=key_id)["Tags"]
+        snapshot.match("list-resource-tags-after-all-untagged", response)
+
+    @markers.aws.validated
+    def test_untag_key_partially(
+        self, kms_client_for_region, kms_create_key, snapshot, region_name
+    ):
+        kms_client = kms_client_for_region(region_name)
+
+        tag_key_to_untag = "tag2"
+        tags = create_tags(**{"tag1": "value1", tag_key_to_untag: "value2", "tag3": "value3"})
+        key_id = kms_create_key(
+            region_name=region_name,
+            Description="test key 123",
+            KeyUsage="ENCRYPT_DECRYPT",
+            Tags=tags,
+        )["KeyId"]
+
+        response = kms_client.list_resource_tags(KeyId=key_id)["Tags"]
+        snapshot.match("list-resource-tags", response)
+
+        kms_client.untag_resource(KeyId=key_id, TagKeys=[tag_key_to_untag])
+
+        response = kms_client.list_resource_tags(KeyId=key_id)["Tags"]
+        snapshot.match("list-resource-tags-after-partially-untagged", response)
+
+    @markers.aws.validated
+    def test_update_and_add_tags_on_tagged_key(
+        self, kms_client_for_region, kms_create_key, snapshot, region_name
+    ):
+        kms_client = kms_client_for_region(region_name)
+
+        tag_key_to_modify = "tag2"
+        tags = create_tags(**{"tag1": "value1", tag_key_to_modify: "value2", "tag3": "value3"})
+        key_id = kms_create_key(
+            region_name=region_name,
+            Description="test key 123",
+            KeyUsage="ENCRYPT_DECRYPT",
+            Tags=tags,
+        )["KeyId"]
+
+        response = kms_client.list_resource_tags(KeyId=key_id)["Tags"]
+        snapshot.match("list-resource-tags", response)
+
+        new_tags = create_tags(
+            **{"tag4": "value4", tag_key_to_modify: "updated_value2", "tag5": "value5"}
+        )
+        kms_client.tag_resource(KeyId=key_id, Tags=new_tags)
+
+        response = kms_client.list_resource_tags(KeyId=key_id)["Tags"]
+        snapshot.match("list-resource-tags-after-tags-updated", response)
+
     @markers.aws.only_localstack
     def test_create_key_custom_id(self, kms_create_key, aws_client):
         custom_id = str(uuid.uuid4())
@@ -987,47 +1084,6 @@ def test_get_put_list_key_policies(self, kms_create_key, aws_client, account_id)
         key_policy = aws_client.kms.get_key_policy(KeyId=key_id, PolicyName="default")["Policy"]
         assert json.dumps(json.loads(key_policy)) == policy_two
 
-    @markers.aws.validated
-    def test_tag_untag_list_tags(self, kms_create_key, aws_client):
-        def _create_tag(key):
-            return {"TagKey": key, "TagValue": short_uid()}
-
-        def _are_tags_there(tags, key_id):
-            if not tags:
-                return True
-            next_token = None
-            while True:
-                kwargs = {"nextToken": next_token} if next_token else {}
-                response = aws_client.kms.list_resource_tags(KeyId=key_id, **kwargs)
-                for response_tag in response["Tags"]:
-                    for i in range(len(tags)):
-                        if response_tag.get("TagKey") == tags[i].get("TagKey") and response_tag.get(
-                            "TagValue"
-                        ) == tags[i].get("TagValue"):
-                            del tags[i]
-                            if not tags:
-                                return True
-                            break
-                if "nextToken" not in response:
-                    break
-                next_token = response["nextToken"]
-            return False
-
-        old_tag_one = _create_tag("one")
-        new_tag_one = _create_tag("one")
-        tag_two = _create_tag("two")
-        tag_three = _create_tag("three")
-
-        key_id = kms_create_key(Tags=[old_tag_one, tag_two])["KeyId"]
-        assert _are_tags_there([old_tag_one, tag_two], key_id) is True
-        # Going to rewrite one of the tags and then add a new one.
-        aws_client.kms.tag_resource(KeyId=key_id, Tags=[new_tag_one, tag_three])
-        assert _are_tags_there([new_tag_one, tag_two, tag_three], key_id) is True
-        assert _are_tags_there([old_tag_one], key_id) is False
-        aws_client.kms.untag_resource(KeyId=key_id, TagKeys=[new_tag_one.get("TagKey")])
-        assert _are_tags_there([tag_two, tag_three], key_id) is True
-        assert _are_tags_there([new_tag_one], key_id) is False
-
     @markers.aws.validated
     def test_cant_use_disabled_or_deleted_keys(self, kms_create_key, aws_client):
         key_id = kms_create_key(KeySpec="SYMMETRIC_DEFAULT", KeyUsage="ENCRYPT_DECRYPT")["KeyId"]

diff --git a/tests/aws/services/kms/test_kms.snapshot.json b/tests/aws/services/kms/test_kms.snapshot.json
@@ -1793,5 +1793,107 @@
         }
       }
     }
+  },
+  "tests/aws/services/kms/test_kms.py::TestKMS::test_tag_existing_key_and_untag": {
+    "recorded-date": "10-01-2025, 09:39:48",
+    "recorded-content": {
+      "list-resource-tags": [
+        {
+          "TagKey": "tag1",
+          "TagValue": "value1"
+        },
+        {
+          "TagKey": "tag2",
+          "TagValue": "value2"
+        }
+      ],
+      "list-resource-tags-after-all-untagged": []
+    }
+  },
+  "tests/aws/services/kms/test_kms.py::TestKMS::test_create_key_with_tag_and_untag": {
+    "recorded-date": "10-01-2025, 09:40:40",
+    "recorded-content": {
+      "list-resource-tags": [
+        {
+          "TagKey": "tag1",
+          "TagValue": "value1"
+        },
+        {
+          "TagKey": "tag2",
+          "TagValue": "value2"
+        }
+      ],
+      "list-resource-tags-after-all-untagged": []
+    }
+  },
+  "tests/aws/services/kms/test_kms.py::TestKMS::test_untag_key_partially": {
+    "recorded-date": "10-01-2025, 09:41:02",
+    "recorded-content": {
+      "list-resource-tags": [
+        {
+          "TagKey": "tag1",
+          "TagValue": "value1"
+        },
+        {
+          "TagKey": "tag2",
+          "TagValue": "value2"
+        },
+        {
+          "TagKey": "tag3",
+          "TagValue": "value3"
+        }
+      ],
+      "list-resource-tags-after-partially-untagged": [
+        {
+          "TagKey": "tag1",
+          "TagValue": "value1"
+        },
+        {
+          "TagKey": "tag3",
+          "TagValue": "value3"
+        }
+      ]
+    }
+  },
+  "tests/aws/services/kms/test_kms.py::TestKMS::test_update_and_add_tags_on_tagged_key": {
+    "recorded-date": "17-01-2025, 12:25:39",
+    "recorded-content": {
+      "list-resource-tags": [
+        {
+          "TagKey": "tag1",
+          "TagValue": "value1"
+        },
+        {
+          "TagKey": "tag2",
+          "TagValue": "value2"
+        },
+        {
+          "TagKey": "tag3",
+          "TagValue": "value3"
+        }
+      ],
+      "list-resource-tags-after-tags-updated": [
+        {
+          "TagKey": "tag1",
+          "TagValue": "value1"
+        },
+        {
+          "TagKey": "tag2",
+          "TagValue": "updated_value2"
+        },
+        {
+          "TagKey": "tag3",
+          "TagValue": "value3"
+        },
+        {
+          "TagKey": "tag4",
+          "TagValue": "value4"
+        },
+        {
+          "TagKey": "tag5",
+          "TagValue": "value5"
+        }
+      ]
+    }
   }
 }
diff --git a/tests/aws/services/kms/test_kms.validation.json b/tests/aws/services/kms/test_kms.validation.json
@@ -23,6 +23,9 @@
   "tests/aws/services/kms/test_kms.py::TestKMS::test_create_key": {
     "last_validated_date": "2024-04-11T15:26:14+00:00"
   },
+  "tests/aws/services/kms/test_kms.py::TestKMS::test_create_key_with_tag_and_untag": {
+    "last_validated_date": "2025-01-10T09:40:40+00:00"
+  },
   "tests/aws/services/kms/test_kms.py::TestKMS::test_create_list_delete_alias": {
     "last_validated_date": "2024-04-11T15:53:50+00:00"
   },
@@ -206,12 +209,21 @@
   "tests/aws/services/kms/test_kms.py::TestKMS::test_sign_verify[RSA_4096-RSASSA_PKCS1_V1_5_SHA_512]": {
     "last_validated_date": "2024-04-11T15:53:06+00:00"
   },
+  "tests/aws/services/kms/test_kms.py::TestKMS::test_tag_existing_key_and_untag": {
+    "last_validated_date": "2025-01-10T09:39:48+00:00"
+  },
   "tests/aws/services/kms/test_kms.py::TestKMS::test_tag_untag_list_tags": {
     "last_validated_date": "2024-04-11T15:53:57+00:00"
   },
+  "tests/aws/services/kms/test_kms.py::TestKMS::test_untag_key_partially": {
+    "last_validated_date": "2025-01-10T09:41:02+00:00"
+  },
   "tests/aws/services/kms/test_kms.py::TestKMS::test_update_alias": {
     "last_validated_date": "2024-04-11T15:53:53+00:00"
   },
+  "tests/aws/services/kms/test_kms.py::TestKMS::test_update_and_add_tags_on_tagged_key": {
+    "last_validated_date": "2025-01-17T12:25:39+00:00"
+  },
   "tests/aws/services/kms/test_kms.py::TestKMS::test_update_key_description": {
     "last_validated_date": "2024-04-11T15:53:46+00:00"
   },