TaggingService: Allow key value field names to be overridden (#12306)

viren-nadkarni · web-flow · commit 6b3878988d8c · 2025-02-28T10:47:50.000+05:30
diff --git a/localstack-core/localstack/utils/tagging.py b/localstack-core/localstack/utils/tagging.py
@@ -2,15 +2,23 @@
 
 
 class TaggingService:
-    def __init__(self):
+    def __init__(self, key_field: str = None, value_field: str = None):
+        """
+        :param key_field: the field name representing the tag key as used by botocore specs
+        :param value_field: the field name representing the tag value as used by botocore specs
+        """
+        self.key_field = key_field or "Key"
+        self.value_field = value_field or "Value"
+
         self.tags = {}
 
     def list_tags_for_resource(self, arn: str, root_name: Optional[str] = None):
         root_name = root_name or "Tags"
+
         result = []
         if arn in self.tags:
             for k, v in self.tags[arn].items():
-                result.append({"Key": k, "Value": v})
+                result.append({self.key_field: k, self.value_field: v})
         return {root_name: result}
 
     def tag_resource(self, arn: str, tags: List[Dict[str, str]]):
@@ -19,7 +27,7 @@ def tag_resource(self, arn: str, tags: List[Dict[str, str]]):
         if arn not in self.tags:
             self.tags[arn] = {}
         for t in tags:
-            self.tags[arn][t["Key"]] = t["Value"]
+            self.tags[arn][t[self.key_field]] = t[self.value_field]
 
     def untag_resource(self, arn: str, tag_names: List[str]):
         tags = self.tags.get(arn, {})
diff --git a/tests/unit/test_tagging.py b/tests/unit/test_tagging.py
@@ -1,30 +1,47 @@
-import unittest
+import pytest
 
 from localstack.utils.tagging import TaggingService
 
 
-class TestTaggingService(unittest.TestCase):
-    svc = TaggingService()
+class TestTaggingService:
+    @pytest.fixture
+    def tagging_service(self):
+        def _factory(**kwargs):
+            return TaggingService(**kwargs)
 
-    def test_list_empty(self):
-        result = self.svc.list_tags_for_resource("test")
-        self.assertEqual({"Tags": []}, result)
+        return _factory
 
-    def test_create_tag(self):
+    def test_list_empty(self, tagging_service):
+        svc = tagging_service()
+        result = svc.list_tags_for_resource("test")
+        assert result == {"Tags": []}
+
+    def test_create_tag(self, tagging_service):
+        svc = tagging_service()
         tags = [{"Key": "key_key", "Value": "value_value"}]
-        self.svc.tag_resource("arn", tags)
-        actual = self.svc.list_tags_for_resource("arn")
+        svc.tag_resource("arn", tags)
+        actual = svc.list_tags_for_resource("arn")
         expected = {"Tags": [{"Key": "key_key", "Value": "value_value"}]}
-        self.assertDictEqual(expected, actual)
+        assert actual == expected
 
-    def test_delete_tag(self):
+    def test_delete_tag(self, tagging_service):
+        svc = tagging_service()
         tags = [{"Key": "key_key", "Value": "value_value"}]
-        self.svc.tag_resource("arn", tags)
-        self.svc.untag_resource("arn", ["key_key"])
-        result = self.svc.list_tags_for_resource("arn")
-        self.assertEqual({"Tags": []}, result)
-
-    def test_list_empty_delete(self):
-        self.svc.untag_resource("arn", ["key_key"])
-        result = self.svc.list_tags_for_resource("arn")
-        self.assertEqual({"Tags": []}, result)
+        svc.tag_resource("arn", tags)
+        svc.untag_resource("arn", ["key_key"])
+        result = svc.list_tags_for_resource("arn")
+        assert result == {"Tags": []}
+
+    def test_list_empty_delete(self, tagging_service):
+        svc = tagging_service()
+        svc.untag_resource("arn", ["key_key"])
+        result = svc.list_tags_for_resource("arn")
+        assert result == {"Tags": []}
+
+    def test_field_name_override(self, tagging_service):
+        svc = tagging_service(key_field="keY", value_field="valuE")
+        tags = [{"keY": "my", "valuE": "congratulations"}]
+        svc.tag_resource("arn", tags)
+        assert svc.list_tags_for_resource("arn") == {
+            "Tags": [{"keY": "my", "valuE": "congratulations"}]
+        }
